Submission #17139897

---

Source Code Expand

use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::Hash;

trait Dfa {
    type Alphabet;
    type State;
    fn init(&self) -> Self::State;
    fn next(&self, s: Self::State, a: Self::Alphabet, i: usize) -> Self::State;
    fn accept(&self, s: Self::State) -> bool;
    fn successful(&self, _: Self::State) -> bool {
        false
    }
    fn unsuccessful(&self, _: Self::State) -> bool {
        false
    }
}

struct And<X, Y>(X, Y);

impl<X: Dfa<Alphabet = A>, Y: Dfa<Alphabet = A>, A: Copy> Dfa for And<X, Y> {
    type Alphabet = A;
    type State = (X::State, Y::State);
    fn init(&self) -> Self::State {
        (self.0.init(), self.1.init())
    }
    fn next(&self, (s0, s1): Self::State, a: Self::Alphabet, i: usize) -> Self::State {
        (self.0.next(s0, a, i), self.1.next(s1, a, i))
    }
    fn accept(&self, (s0, s1): Self::State) -> bool {
        self.0.accept(s0) && self.1.accept(s1)
    }
    fn successful(&self, (s0, s1): Self::State) -> bool {
        self.0.successful(s0) && self.1.successful(s1)
    }
    fn unsuccessful(&self, (s0, s1): Self::State) -> bool {
        self.0.unsuccessful(s0) || self.1.unsuccessful(s1)
    }
}

struct Leq<'a>(&'a [u8]);

impl Dfa for Leq<'_> {
    type Alphabet = u8;
    type State = Ordering;
    fn init(&self) -> Self::State {
        Ordering::Equal
    }
    // assumes i moves from 0 to self.0.len() - 1
    fn next(&self, s: Self::State, a: Self::Alphabet, i: usize) -> Self::State {
        s.then(a.cmp(&self.0[i]))
    }
    fn accept(&self, s: Self::State) -> bool {
        s != Ordering::Greater
    }
    fn successful(&self, s: Self::State) -> bool {
        s == Ordering::Less
    }
    fn unsuccessful(&self, s: Self::State) -> bool {
        s == Ordering::Greater
    }
}

struct DigitSumMultipleOf(u32);

impl Dfa for DigitSumMultipleOf {
    type Alphabet = u8;
    type State = u32;
    fn init(&self) -> Self::State {
        0
    }
    fn next(&self, s: Self::State, a: Self::Alphabet, _: usize) -> Self::State {
        (s + (a - b'0') as u32) % self.0
    }
    fn accept(&self, s: Self::State) -> bool {
        s == 0
    }
}

fn count<X: Dfa>(dfa: &X, n: usize, modulo: u32, alphabet: &[X::Alphabet]) -> u32
where
    X::Alphabet: Copy,
    X::State: Eq + Hash + Copy,
{
    let mut dp = HashMap::new();
    let mut dp2 = HashMap::new();
    dp.insert(dfa.init(), 1_u64);
    for i in 0..n {
        for (s, k) in dp.drain() {
            let k = k % modulo as u64;
            for &a in alphabet {
                let s1 = dfa.next(s, a, i);
                if dfa.unsuccessful(s1) {
                    continue;
                }
                *dp2.entry(s1).or_insert(0) += k;
            }
        }
        std::mem::swap(&mut dp, &mut dp2);
    }
    let mut sum = 0;
    for (s, k) in dp {
        if dfa.accept(s) {
            sum += k;
            sum %= modulo as u64
        }
    }
    sum as u32
}

use proconio::{input, marker::Bytes};

fn main() {
    input! {
        d: u32,
        n: Bytes,
    }
    let dfa = And(Leq(&n), DigitSumMultipleOf(d));
    let alphabet = "0123456789".as_bytes();
    let ans = count(&dfa, n.len(), 1_000_000_007, alphabet);
    println!("{}", ans - 1);
}

Submission Info

Judge Result