Typical DP コンテストの桁DP

https://atcoder.jp/contests/tdpc/tasks/tdpc_number

Untitled

↓実装

impl Dp {
    fn new(n_digit: usize, d: i64) -> Self {
        Self {
            dp: vec![vec![vec![Mint::new(0); d as usize]; 2]; n_digit + 1],
        }
    }

    fn at(&self, digit_i: usize, smaller: bool, rem: i64) -> &Mint {
        &self.dp[digit_i][smaller as usize][rem as usize]
    }

    fn at_mut(&mut self, digit_i: usize, smaller: bool, rem: i64) -> &mut Mint {
        &mut self.dp[digit_i][smaller as usize][rem as usize]
    }
}
impl Problem {
    fn solve(&self) -> Answer {
        let mut dp = Dp::new(self.k.len(), self.d);

        *dp.at_mut(0, false, 0) = 1.into();
        for digit_i in 0..self.k.len() {
            for rem in 0..self.d {
                let dp_false = *dp.at(digit_i, false, rem);
                let dp_true = *dp.at(digit_i, true, rem);

                for x in 0..10 {
                    *dp.at_mut(digit_i + 1, true, (rem + x) % self.d) += dp_true;
                }

                for x in 0..self.k[digit_i] {
                    *dp.at_mut(digit_i + 1, true, (rem + x) % self.d) += dp_false;
                }

                *dp.at_mut(digit_i + 1, false, (rem + self.k[digit_i]) % self.d) += dp_false;
            }
        }
        // 0 を除くために 1 を引く
        let ans = dp.at(self.k.len(), true, 0) + dp.at(self.k.len(), false, 0) - Mint::new(1);
        let ans = ans.val() as i64;
        Answer { ans }
    }
}

その他