AtCoder Beginner Contest 362 E

https://atcoder.jp/contests/abc362/tasks/abc362_e

DPです。状態は(次の値, 公差, 長さ)、値は頻度にすればよいですが、次の値ですぐに検索できるようにHashMapのキーを次の値にしておきます。

// Count Arithmetic Subsequences
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// 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 ////////////////////

const D: usize = 998244353;

fn read_input() -> Vec<i64> {
    let _N: usize = read();
    let A: Vec<i64> = read_vec();
    A
}

fn F(A: Vec<i64>) {
    let N = A.len();
    // { next: { (d, len): freq } }
    let mut dp: HashMap<i64, HashMap<(i64, usize), usize>> = HashMap::new();
    for i in 1..N {
        let mut v: Vec<(i64, usize, usize)> = vec![];   // [(d, len, freq)]
        if let Some(dic) = dp.get(&A[i]) {
            for (&(d1, len1), &freq1) in dic.iter() {
                v.push((d1, len1+1, freq1))
            }
        }
        for j in 0..i {
            v.push((A[i] - A[j], 2, 1))
        }
        
        for (d, len, freq) in v.into_iter() {
            let n = A[i] + d;
            let e1 = dp.entry(n).or_insert(HashMap::new());
            let e2 = (*e1).entry((d, len)).or_insert(0);
            *e2 = (*e2 + freq) % D;
        }
    }
    
    let mut B: Vec<usize> = vec![0; N+1];
    B[1] = N;
    for (_, dic) in dp.into_iter() {
        for ((_, len), freq) in dic.into_iter() {
            B[len] = (B[len] + freq) % D
        }
    }
    
    println!("{}", (1..N+1).map(|i| B[i].to_string()).
                                    collect::<Vec<String>>().join(" "))
}

fn main() {
    let A = read_input();
    F(A)
}