AtCoder Beginner Contest 438 G

https://atcoder.jp/contests/abc438/tasks/abc438_g

これも足し算の順番を変えて解きます。順番を変えるというのは、どこを固定するかという話です。ここではBを固定します。入力例1だとまずBの最初の1とそれに対応するAを考えます。Bのindex:i mod Mが0ときにAのindex:i mod Nは0, 2, 1となります。具体的には、3, 4, 1です。このときBは1でこれより小さいAはありません。なので、3*1 = 3となります。個々のBに対応するAで木を作って、Bより小さい値だけ和を取ります。そして、足りない個数ぶんだけBを足します。
次のBも同様に考えます。そうすると、今度はAのindexは1, 0となります。さっきAのindexが0, 2, 1でしたが、最初にMより小さいindexを調べると、0と2を削除して0を追加すればよいです。このように木を少ない計算量で更新できます。
もうちょっと具体的に考えて、N=4, M=3, K=20で考えると、最初のBに対応するAは7つあって、0, 3, 2, 1, 0, 3, 2となります。次にMより小さいのindexは2なので、Bのindex:2に対応するAのindexは2つ削ってトータルの個数は1つ減るから1つ足して、2, 1, 0, 3, 2, 1となります。次のMより小さいindexは1なので、1つ削ってトータルの個数はまた1つ増えるので2つ増やして1, 0, 3, 2, 1, 0, 3となります。
このアルゴリズムはAがBより短いと成り立たないので、そのときはAとBを入れ替えます。

// Sum of Min
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}

fn gcd(a: usize, b: usize) -> usize {
    if b == 0 { a } else { gcd(b, a%b) }
}


//////////////////// SegTree ////////////////////

trait Monoid {
    type S: Clone;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn e() -> Self::S;
}

struct SegTree<M: Monoid> {
    n: usize,
    size: usize,
    data: Vec<M::S>,
}

impl<M: Monoid> SegTree<M> {
    fn new(n: usize) -> Self {
        let size = n.next_power_of_two();
        SegTree {
            n,
            size,
            data: vec![M::e(); size*2-1],
        }
    }
    
    fn from(v: Vec<M::S>) -> Self {
        let n = v.len();
        let size = n.next_power_of_two();
        let mut data = vec![M::e(); size*2-1];
        for i in 0..n {
            data[i+size-1] = v[i].clone();
        }
        let mut seg = SegTree { n, size, data };
        for i in (0..size-1).rev() {
            seg.data[i] = M::op(&seg.data[2*i+1], &seg.data[2*i+2]);
        }
        seg
    }
    
    fn set(&mut self, mut i: usize, x: M::S) {
        i = i + self.size - 1;
        self.data[i] = x;
        while i != 0 {
            i = (i - 1) / 2;
            self.data[i] = M::op(&self.data[2*i+1], &self.data[2*i+2]);
        }
    }
    
    fn prod(&self, mut l: usize, mut r: usize) -> M::S {
        let mut sml = M::e();
        let mut smr = M::e();
        l = l + self.size - 1;
        r = r + self.size - 1;
        while l < r {
            if l % 2 == 0 {
                sml = M::op(&sml, &self.data[l]);
                l += 1;
            }
            if r % 2 == 0 {
                r -= 1;
                smr = M::op(&self.data[r], &smr);
            }
            l = (l - 1) / 2;
            r = (r - 1) / 2
        }
        M::op(&sml, &smr)
    }
    
    fn get(&self, i: usize) -> M::S {
        self.data[i+self.size-1].clone()
    }
}


//////////////////// Monoid ////////////////////

struct Sum;
impl Monoid for Sum {
    type S = (usize, usize);    // (個数, 和)
    fn op((a1, b1): &Self::S, (a2, b2): &Self::S) -> Self::S {
        ((a1+a2)%D, (b1+b2)%D)
    }
    fn e() -> Self::S { (0, 0) }
}


//////////////////// Tree ////////////////////

const D: usize = 998244353;

struct Tree {
    seg: SegTree::<Sum>,
    ids: Vec<usize>,
    A: Vec<usize>,
    A_sorted: Vec<usize>
}

impl Tree {
    fn size(&self) -> usize {
        self.seg.size
    }
    
    fn increment(&mut self, i: usize) {
        let id = self.ids[i];
        let (n, v) = self.seg.get(id);
        self.seg.set(id, (n+1, (v+self.A[i]) % D))
    }
    
    fn decrement(&mut self, i: usize) {
        let id = self.ids[i];
        let (n, v) = self.seg.get(id);
        self.seg.set(id, (n-1, (v+D+D-self.A[i]) % D))
    }
    
    // 与えられた値以上のself.A_sortedの最初のindex
    fn first_index(&self, a: usize) -> usize {
        let N = self.A_sorted.len();
        let mut first: usize = 0;
        let mut last: usize = N;
        while last - first > 1 {
            let mid = (first + last) / 2;
            if self.A_sorted[mid] > a {
                last = mid
            }
            else {
                first = mid
            }
        }
        if self.A_sorted[first] >= a {
            first
        }
        else {
            first + 1
        }
    }
    
    // 与えられた値未満の総和
    fn sum(&self, a: usize) -> (usize, usize) {
        let index = self.first_index(a);
        self.seg.prod(0, index)
    }
    
    // Aをソートして重複をなくした配列と
    // Aのidからその配列のidへの配列を作る
    fn make_sorted_id_table(A: &Vec<usize>) -> (Vec<usize>, Vec<usize>) {
        // A:       [7, 7, 3, 5]
        // A_sorted:[3, 5, 7]
        // ids:     [2, 2, 0, 1]
        let N = A.len();
        let mut B: Vec<(usize, usize)> = A.iter().enumerate().
                                            map(|(i, &a)| (a, i)).collect();
        B.sort();
        let mut ids: Vec<usize> = vec![0; N];
        let mut k: usize = 0;
        let mut prev_a: usize = B[0].0;
        let mut C: Vec<usize> = vec![prev_a];
        for &(a, i) in B.iter() {
            if a != prev_a {
                prev_a = a;
                C.push(a);
                k += 1
            }
            ids[i] = k
        }
        (C, ids)
    }
    
    fn create(K: usize, M: usize, A: Vec<usize>) -> Self {
        let N = A.len();
        let (A_sorted, ids) = Tree::make_sorted_id_table(&A);
        let mut seg = SegTree::<Sum>::new(A_sorted.len());
        let Q = (K + M - 1) / M;    // B0が現れる回数
        for i in 0..N {
            let j = i * M % N;
            let q = (Q + N - i - 1) / N;    // Aが現れる回数
            let (n, s) = seg.get(ids[j]);
            seg.set(ids[j], ((n+q)%D, (s+q%D*A[j])%D))
        }
        Tree { seg, ids, A, A_sorted }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<usize>, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let K = v[2];
    let A: Vec<usize> = read_vec();
    let B: Vec<usize> = read_vec();
    (K, A, B)
}

fn F_each(K: usize, A: Vec<usize>, B: Vec<usize>) -> usize {
    let N = A.len();
    let M = B.len();
    let q = K / M;
    let r = K % M;
    let mut tree = Tree::create(K, M, A);
    let (n1, s1) = tree.sum(B[0]);
    let n0 = (K + M - 1) / M;
    let mut s = (B[0] * ((n0 - n1) % D) + s1) % D;
    
    let mut j = M % N;  // Aのindex
    let mut jump: usize = 1;
    let mut prev_j = 0;
    while j != 0 {
        if j >= M {
            jump += 1;
            j = (j + M) % N;
            continue
        }
        
        // 同じBがいくつ出てくるか
        let prev_num = if prev_j < r { q+1 } else { q };
        let num = if j < r { q+1 } else { q };
        // 何個追加するか
        let num_additional = jump + num - prev_num;
        // 追加する
        for k in 0..num_additional {
            let j1 = (prev_j+M*(prev_num+k))%N;
            tree.increment(j1)
        }
        
        // 取り除く
        for k in 0..jump {
            let j2 = (prev_j+M*k)%N;
            tree.decrement(j2)
        }
        let (n2, s2) = tree.sum(B[j]);
        s = (s + B[j] * ((num - n2) % D) + s2) % D;
        prev_j = j;
        j = (prev_j + M) % N;
        jump = 1
    }
    s
}

fn F(K: usize, A: Vec<usize>, B: Vec<usize>) -> usize {
    let N = A.len();
    let M = B.len();
    if N < M {
        return F(K, B, A)
    }
    
    let mut s: usize = 0;
    let d = gcd(N, M);
    for i in 0..d {
        let A1: Vec<usize> = (i..N).step_by(d).map(|j| A[j]).collect();
        let B1: Vec<usize> = (i..M).step_by(d).map(|j| B[j]).collect();
        s = (s + F_each((K+d-1-i)/d, A1, B1)) % D
    }
    s
}

fn main() {
    let (K, A, B) = read_input();
    println!("{}", F(K, A, B))
}