K-Divisible Subarrays の解説

例えば、Ai, Ai+1, Ai+2, Ai+3で条件を満たせつつ、 Ai+1, Ai+2でも満たせるならば後者を選んだ方が良いので 満たせるものができたら貪欲に左から作る方が良い。

ということで、 dp[i] := Aiを必ず使う前提でiまでで作れるmod k == 0の区間の個数 をセグ木で持っておくことで解くことができる。

iを固定して考えた時に、iを使ってmod k == 0を満たせる区間の一番右のindexをmapとして持っておくことでそれより小さいindexのsegのクエリ + 1をdp[i]として更新できる。

fn main() {
    input! {
        n: usize, k: usize,
        a: [usize; n]
    }
    let mut map = HashMap::new();
    let mut dp = Segtree::<Max<usize>>::new(n);
    map.insert(0, 0);

    let mut acc = 0;
    for i in 0..n {
        let num = a[i];
        let target = (acc + num) % k;
        if let Some(&v) = map.get(&target) {
            let pre = dp.prod(..v);
            dp.set(i, pre + 1);
        }
        acc += num;
        acc %= k;
        map.insert(acc, i + 1);
    }
    let ans = dp.all_prod();
    println!("{}" ,ans);
}