AtCoder Beginner Contest 336 E

https://atcoder.jp/contests/abc336/tasks/abc336_e

ある桁和を考えて、上の桁から、数字の和と剰余を状態としたDPを考えればよいです。上限があるのでちょっと嫌らしいです。

// Digit Sum Divisible
#![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 digits(N: i64) -> Vec<i64> {
    let mut ds: Vec<i64> = vec![];
    let mut n = N;
    while n != 0 {
        let d = n % 10;
        ds.push(d);
        n /= 10
    }
    ds
}


//////////////////// DP ////////////////////

use std::collections::HashMap;
type Map = HashMap<(i64, i64), i64>;

struct DP {
    dp1: Map,
    dp2: Map,
    d: i64
}

impl DP {
    fn update(&self, d1: i64) -> DP {
        let mut dp1: Map = HashMap::new();
        let mut dp2: Map = HashMap::new();
        self.update_core(d1, d1 + 1, &self.dp1, &mut dp1);
        self.update_core(0, d1, &self.dp1, &mut dp2);
        self.update_core(0, 10, &self.dp2, &mut dp2);
        DP { dp1, dp2, d: self.d }
    }
    
    fn update_core(&self, first: i64, last: i64,
                                dp_in: &Map, dp_out: &mut Map) {
        for (&(s, r), &n) in dp_in.iter() {
            for d in first..last {
                let key = (s + d, (r * 10 + d) % self.d);
                let e = dp_out.entry(key).or_insert(0);
                *e += n
            }
        }
    }
    
    fn count_core(&self, dp: &Map) -> i64 {
        let mut counter: i64 = 0;
        for (&(s, r), &n) in dp.iter() {
            if s == self.d && r == 0 {
                counter += n
            }
        }
        counter
    }
    
    fn count(&self) -> i64 {
        self.count_core(&self.dp1) + self.count_core(&self.dp2)
    }
    
    fn new(d: i64) -> DP {
        let mut dp1: Map = HashMap::new();
        dp1.insert((0, 0), 1);
        let dp2: Map = HashMap::new();
        DP { dp1, dp2, d }
    }
}


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

fn count_each(d: i64, ds: &Vec<i64>) -> i64 {
    let mut dp = DP::new(d);
    for &d1 in ds.iter().rev() {
        dp = dp.update(d1)
    }
    dp.count()
}

fn F(N: i64) -> i64 {
    let ds = digits(N);
    let max_sum = ds[ds.len()-1] + (ds.len() as i64 - 1) * 9;
    let mut counter: i64 = 0;
    for d in 1..max_sum+1 {
        counter += count_each(d, &ds)
    }
    counter
}

fn main() {
    let N = read();
    println!("{}", F(N))
}