この記事は 夏のブログリレー 10 日目の記事です。

こんにちは、23B の @hayatroid です。院試対策でオートマトンの勉強をしたので、同じく講義などでオートマトンの勉強をした人間を競プロに誘いです。

TL;DR

kuretchi さんのオートマトン上の DP に KMP オートマトンと Aho–Corasick オートマトンを impl です。そして次の問題を AC です。

• yukicoder No.220 世界のなんとか2
• yukicoder No.260 世界のなんとか3
• yukicoder No.2867 NOT FOUND 404 Again
• yukicoder No.1269 I hate Fibonacci Number

ナベアツとはオートマトンである

ナベアツ^[1]は、$3$ の倍数と $3$ がつくときアホになる。数を読んでアホになるかどうかを決めるのだから、ナベアツはオートマトンといえる。

ナベアツ数を認識するオートマトン

非負整数がナベアツ数であるとは、$3$ の倍数であるか $3$ がつくことをいう^[2]。

$3$ の倍数を認識するオートマトンと $3$ がつく数を認識するオートマトンは、それぞれ次のとおりである。

$3$ の倍数	$3$ がつく

ナベアツ数を認識するオートマトンは、この $2$ つの union で表せる。

たとえば $10125$ は、$3$ がつかないが $3$ の倍数であるため、ナベアツ数である。

$3$ の倍数 ($10125$ を読む)	$3$ がつく ($10125$ を読む)

たとえば $38347$ は、$3$ の倍数でないが $3$ がつくため、ナベアツ数である。

$3$ の倍数 ($38347$ を読む)	$3$ がつく ($38347$ を読む)

これらのオートマトンを実装すれば、次の問題が解ける。

No.220 世界のなんとか2

yukicoder No.220 世界のなんとか2 では、$1$ 以上 $10^P$ 以下のナベアツ数を数え上げる必要がある。$P$ は最大 $19$ である。

先の図のオートマトンを実装し、この問題を解く。

まず、決定性有限オートマトン^[3]の $5$ つ組をそのまま trait にする。

trait Dfa {
    type State;
    type Alphabet;
    fn init(&self) -> Self::State;
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State;
    fn accept(&self, q: &Self::State) -> bool;
}

先の図の左、$3$ の倍数を認識するオートマトンを MultipleOf(3) として実装する。

struct MultipleOf(u64);
impl Dfa for MultipleOf {
    type State = u64;
    type Alphabet = u8;
    fn init(&self) -> Self::State {
        0
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        (q * 10 + (c - b'0') as u64) % self.0
    }
    fn accept(&self, q: &Self::State) -> bool {
        *q == 0
    }
}

先の図の右、$3$ がつく数を認識するオートマトンを Seen(b'3') として実装する。

struct Seen(u8);
impl Dfa for Seen {
    type State = bool;
    type Alphabet = u8;
    fn init(&self) -> Self::State {
        false
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        *q || *c == self.0
    }
    fn accept(&self, q: &Self::State) -> bool {
        *q
    }
}

オートマトンの union もまたオートマトンである。$2$ つのオートマトンを並べて走らせ、どちらかが受理すれば受理するオートマトンを Or として実装する。

struct Or<A, B>(A, B);
impl<A: Dfa, B: Dfa<Alphabet = A::Alphabet>> Dfa for Or<A, B> {
    type State = (A::State, B::State);
    type Alphabet = A::Alphabet;
    fn init(&self) -> Self::State {
        (self.0.init(), self.1.init())
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        (self.0.next(&q.0, c), self.1.next(&q.1, c))
    }
    fn accept(&self, q: &Self::State) -> bool {
        self.0.accept(&q.0) || self.1.accept(&q.1)
    }
}

オートマトン $A$ に ${\Sigma }^n$ のすべての文字列を読ませ、$A$ が受理する文字列を数える関数を $\mathrm{count}(A,\Sigma ,n)$ として実装する。たとえば $\Sigma$ を $10$ 進の数字とすれば、$\mathrm{count}(A,\Sigma ,n)$ は $n$ 桁の数字列のうち $A$ が受理するものを数える。先頭の $0$ を気にしないことにすれば、これは $10^n$ 未満の非負整数のうち $A$ が受理するものの個数とみなせる。

fn count<A, S>(dfa: A, sigma: S, len: usize) -> u64
where
    A: Dfa,
    A::State: Eq + Hash,
    S: Iterator<Item = A::Alphabet> + Clone,
{
    let mut dp = HashMap::new();
    dp.insert(dfa.init(), 1);
    for _ in 0..len {
        let mut ndp = HashMap::new();
        for (q, v) in dp {
            for c in sigma.clone() {
                *ndp.entry(dfa.next(&q, &c)).or_insert(0) += v;
            }
        }
        dp = ndp;
    }
    dp.iter()
        .filter_map(|(q, v)| dfa.accept(q).then_some(v))
        .sum()
}

あとはオートマトンを組み立てて長さ $P$ で数え上げるのみである。$10^P$ はナベアツ数でなく、$0$ はナベアツ数なので、$1$ を引けば答えになる。

fn main() {
    input!(p: usize);
    let dfa = Or(MultipleOf(3), Seen(b'3'));
    println!("{}", count(dfa, b'0'..=b'9', p) - 1);
}

以上を提出すると AC が得られる。

No.260 世界のなんとか3

実は先の問題はナベアツ方程式^[4]で解けてしまう。それでも問題をオートマトンで解くことの真髄は、条件が増えてもオートマトンを足すだけで済むことにある。

yukicoder No.260 世界のなんとか3 では、$A$ 以上 $B$ 以下の整数のうち、ナベアツ数であって $8$ の倍数でないものを $10^9+7$ で割った余りで数え上げる必要がある。$A,B$ は最大 $10^{10000}$ である。

これは、先の問題で用意したオートマトンに加えて、次の $4$ つのオートマトンを足せば解ける。

• $B$ 以下を認識するオートマトン Le
• $A$ 未満を認識するオートマトン Lt
• オートマトンの intersection And
• オートマトンの complement Not

たとえば $B=31415$ のとき、$B$ 以下を認識するオートマトンは次のとおりである。読んだ数字列が $31415$ の接頭辞より小さいか、等しいか、大きいかを状態に持つ。

$31415$ 以下

上図のオートマトンを Le(b"31415") として実装する。

struct Le<'a>(&'a [u8]);
impl Dfa for Le<'_> {
    type State = (Ordering, usize);
    type Alphabet = u8;
    fn init(&self) -> Self::State {
        (Ordering::Equal, 0)
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        (q.0.then(c.cmp(&self.0[q.1])), q.1 + 1)
    }
    fn accept(&self, q: &Self::State) -> bool {
        q.0.is_le()
    }
}

Lt, And, Not もこれまで同様のノリで実装する。count は $10^9+7$ で割った余りを取るよう書き換える。

あとはオートマトンを組み立てて、$B$ 以下の数え上げから $A$ 未満の数え上げを引くのみである。

fn main() {
    input!(a: Bytes, b: Bytes);
    let nabeatsu = || And(Or(MultipleOf(3), Seen(b'3')), Not(MultipleOf(8)));
    let dfa_le_b = And(Le(&b), nabeatsu());
    let dfa_lt_a = And(Lt(&a), nabeatsu());
    let le_b = count(dfa_le_b, b'0'..=b'9', b.len());
    let lt_a = count(dfa_lt_a, b'0'..=b'9', a.len());
    println!("{}", (le_b + MOD - lt_a) % MOD);
}

以上を提出すると AC が得られる。

ナベアツが本気を出してきたらどうしよう

今度のナベアツは、$3$ の倍数と $w$ がつくときアホになるという。

$w$ がつく数を認識するオートマトン

$w$ がつく数を認識するオートマトンは、読んだ数字列の末尾と $w$ の先頭が一致する長さのうち、最も長いものを状態に持てばよい。読んだ数字が $w$ の次の文字と一致しなかったら、その次に長いものへ戻り、そこから続きを読む。この戻る辺を failure link という^[5]。$w$ を一列に並べて failure link を足したものを KMP オートマトンという^[5:1]。

たとえば $w=33$ のとき、$w$ がつく数を認識するオートマトンは次のとおりである。破線が failure link である。読む数字の辺が無いときは、破線を $1$ 本たどって読む。それでも読めなければ、また破線を $1$ 本たどって読む。$0$ まで戻っても読めなければ、$0$ に留まる。

$33$ がつく

たとえば $344333$ は、$3$ の倍数でないが $33$ がつくため、ナベアツ数である。

$33$ がつく ($344333$ を読む)

このオートマトンを実装すれば、次の問題が解ける。

No.2867 NOT FOUND 404 Again

yukicoder No.2867 NOT FOUND 404 Again では、$N$ 以下の正の整数のうち $404$ がつかないものを $998244353$ で割った余りで数え上げる必要がある。$N$ は最大 $10^{1000000}$ である。

これは、先の問題のオートマトンに加えて、$w$ がつく数を認識するオートマトン Contains を足せば解ける。

先の図のオートマトンを Contains::new(b"33") として実装する。

状態 $q$ において、数字 $c$ を読んだ次の状態を next[q][c]、破線を $1$ 本たどった先を fail[q] とする。$c$ の辺が無いときは、破線をたどって読んだ先を next[q][c] とする。next と fail は、$q$ の小さいほうから埋められる。

struct Contains {
    next: Vec<[usize; 10]>,
    accept: Vec<bool>,
}
impl Contains {
    fn new(w: &[u8]) -> Self {
        // 1. w を一列に並べる
        let n = w.len();
        let mut next = vec![[0; 10]; n + 1];
        let mut accept = vec![false; n + 1];
        for q in 0..n {
            next[q][(w[q] - b'0') as usize] = q + 1;
        }
        accept[n] = true;
        // 2. failure link を張る
        let mut fail = vec![0; n + 1];
        for q in 1..=n {
            for c in 0..10 {
                if next[q][c] != 0 {
                    fail[next[q][c]] = next[fail[q]][c];
                } else {
                    next[q][c] = next[fail[q]][c];
                }
            }
        }
        // 3. 受理状態に留まるようにする
        next[n] = [n; 10];
        Contains { next, accept }
    }
}
impl Dfa for Contains {
    type State = usize;
    type Alphabet = u8;
    fn init(&self) -> Self::State {
        0
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        self.next[*q][(c - b'0') as usize]
    }
    fn accept(&self, q: &Self::State) -> bool {
        self.accept[*q]
    }
}

count は $998244353$ で割った余りを取るよう書き換える。

あとはオートマトンを組み立てて長さ $|N|$ で数え上げるのみである。$0$ は $404$ がつかないので、$1$ を引けば答えになる。

fn main() {
    input!(n: Bytes);
    let dfa = And(Not(Contains::new(b"404")), Le(&n));
    println!("{}", (count(dfa, b'0'..=b'9', n.len()) + MOD - 1) % MOD);
}

以上を提出すると AC が得られる。

Re: ナベアツが本気を出してきたらどうしよう

今度のナベアツは、$3$ の倍数と $W$ のどれかがつくときアホになるという。

$W$ のどれかがつく数を認識するオートマトン

$W$ のどれかがつく数を認識するオートマトンは、KMP オートマトンの一列を $W$ の trie^[6] に差し替え、同じように failure link を足したものである。これを Aho–Corasick オートマトン という^[7]。

たとえば $W=\{33,313\}$ のとき、$W$ のどれかがつく数を認識するオートマトンは次のとおりである。破線が failure link である。受理状態に添えた集合が、そこで見つかる語である。見つかる語は、破線を $1$ 本たどった先で見つかる語を含む。その先の語も、また破線を $1$ 本たどった先の語を含む。

$33$ または $313$ がつく

たとえば $31131335533$ は、$3$ の倍数でないが $313$ も $33$ もつくため、ナベアツ数である。

$33$ または $313$ がつく ($31131335533$ を読む)

このオートマトンを実装すれば、次の問題が解ける。

No.1269 I hate Fibonacci Number

yukicoder No.1269 I hate Fibonacci Number では、$10^N$ 未満の正の整数のうち、$L$ 以上 $R$ 以下のフィボナッチ数がどれもつかないものを $10^9+7$ で割った余りで数え上げる必要がある。$N$ は最大 $5000$、$R$ は最大 $10^{18}$ である。

これは、先の問題のオートマトンに加えて、$W$ のどれかがつく数を認識するオートマトン ContainsAny を足せば解ける。

先の図のオートマトンを ContainsAny::new(&[b"33".to_vec(), b"313".to_vec()]) として実装する。

状態 $q$ において、数字 $c$ を読んだ次の状態を next[q][c]、破線を $1$ 本たどった先を fail[q] とする。$c$ の辺が無いときは、破線をたどって読んだ先を next[q][c] とする。next と fail は、根からの深さが小さいほうから埋められ、幅優先探索をすればよい。実装は OI Wiki の AC 自动机 を参考にした。

struct ContainsAny {
    next: Vec<[usize; 10]>,
    accept: Vec<bool>,
}
impl ContainsAny {
    fn new(ws: &[Vec<u8>]) -> Self {
        // 1. ws の trie をつくる
        let mut next = vec![[0; 10]];
        let mut accept = vec![false];
        for w in ws {
            let mut q = 0;
            for c in w {
                let c = (c - b'0') as usize;
                if next[q][c] == 0 {
                    next[q][c] = next.len();
                    next.push([0; 10]);
                    accept.push(false);
                }
                q = next[q][c];
            }
            accept[q] = true;
        }
        // 2. failure link を張る
        let mut fail = vec![0; next.len()];
        let mut bfs = VecDeque::new();
        for c in 0..10 {
            if next[0][c] != 0 {
                bfs.push_back(next[0][c]);
            }
        }
        while let Some(q) = bfs.pop_front() {
            for c in 0..10 {
                if next[q][c] != 0 {
                    fail[next[q][c]] = next[fail[q]][c];
                    bfs.push_back(next[q][c]);
                } else {
                    next[q][c] = next[fail[q]][c];
                }
            }
            accept[q] |= accept[fail[q]];
        }
        // 3. 受理状態に留まるようにする
        for q in 0..accept.len() {
            if accept[q] {
                next[q] = [q; 10];
            }
        }
        ContainsAny { next, accept }
    }
}
impl Dfa for ContainsAny {
    type State = usize;
    type Alphabet = u8;
    fn init(&self) -> Self::State {
        0
    }
    fn next(&self, q: &Self::State, c: &Self::Alphabet) -> Self::State {
        self.next[*q][(c - b'0') as usize]
    }
    fn accept(&self, q: &Self::State) -> bool {
        self.accept[*q]
    }
}

count は $10^9+7$ で割った余りを取るよう書き換える。

あとはフィボナッチ数を列挙し、オートマトンを組み立てて長さ $N$ で数え上げるのみである。$0$ はフィボナッチ数がつかないので、$1$ を引けば答えになる。

fn main() {
    input!(n: usize, l: u64, r: u64);
    let mut ws = vec![];
    let (mut x, mut y) = (1, 1);
    while x < l {
        (x, y) = (y, x + y);
    }
    while x <= r {
        ws.push(x.to_string().into_bytes());
        (x, y) = (y, x + y);
    }
    let dfa = Not(ContainsAny::new(&ws));
    println!("{}", (count(dfa, b'0'..=b'9', n) + MOD - 1) % MOD);
}

以上を提出すると AC が得られる。

〆

Rust の trait があると、オートマトンの具体を与えずに、抽象からオートマトンを組み立てたり (Or)、抽象のまま数え上げたり (count) できて面白いです。

以上、オートマトン er を競プロに誘い記事と見せかけた、競プロ er を Rust に誘い記事でした。