AtCoder Beginner Contest 427 F

https://atcoder.jp/contests/abc427/tasks/abc427_f

長さ60ですが、半分だと2³⁰通りの部分列があるわけではなく、隣は空けないといけないのでフィボナッチ数になって200万程度しかないので全部を部分列でMの剰余がいくつになるかを求められます。前半と後半で求めてソートすればすぐに剰余0の組合せの数を求められますが、隣がないのを考慮するために後半は反対向きに求めます。

// Not Adjacent
#![allow(non_snake_case)]

use std::collections::BTreeMap;


//////////////////// 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()
}


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

fn read_input() -> (i32, Vec<i32>) {
    let v: Vec<i32> = read_vec();
    let M = v[1];
    let A: Vec<i32> = read_vec();
    (M, A)
}

fn extend(v: BTreeMap<i32, usize>, w: &mut BTreeMap<i32, usize>) {
    for (m, f) in v {
        let e = w.entry(m).or_insert(0);
        *e += f
    }
}

fn collect_sums(B: &[i32], M: i32) -> (Vec<(i32, usize)>, Vec<(i32, usize)>) {
    let mut v: BTreeMap<i32, usize> = BTreeMap::new();  // 直近の要素を取った
    let mut w: BTreeMap<i32, usize> = BTreeMap::new();  // それ以外
    v.insert(B[0], 1);
    w.insert(0, 1);
    for &n in &B[1..] {
        let new_v: BTreeMap<i32, usize> = w.iter().map(|(&m, &f)| ((m+n)%M, f)).
                                                                    collect();
        extend(v, &mut w);
        v = new_v;
    }
    (v.into_iter().collect::<Vec<(i32, usize)>>(),
     w.into_iter().collect::<Vec<(i32, usize)>>())
}

fn count_multiples(v1: &Vec<(i32, usize)>,
                    v2: &Vec<(i32, usize)>, M: i32) -> usize {
    let mut counter: usize = 0;
    let mut k = 0;
    let mut l = 0;
    while k < v1.len() && l < v2.len() {
        if v1[k].0 + v2[v2.len()-1-l].0 == M {
            counter += v1[k].1 * v2[v2.len()-1-l].1;
            k += 1;
            l += 1;
        }
        else if v1[k].0 + v2[v2.len()-1-l].0 < M {
            k += 1
        }
        else {
            l += 1
        }
    }
    
    if v1[0].0 == 0 && v2[0].0 == 0 {
        counter += v1[0].1 * v2[0].1
    }
    counter
}

fn F(M: i32, A: Vec<i32>) -> usize {
    if A.len() == 1 {
        return if A[0] % M == 0 { 2 } else { 1 }
    }
    
    let mid = A.len() / 2;
    let (v1, w1) = collect_sums(&A[..mid], M);
    let B: Vec<i32> = (mid..A.len()).rev().map(|i| A[i]).collect();
    let (v2, w2) = collect_sums(&B[..], M);
    let n1 = count_multiples(&v1, &w2, M);
    let n2 = count_multiples(&w1, &v2, M);
    let n3 = count_multiples(&w1, &w2, M);
    n1 + n2 + n3
}

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