z_algorithm/
lib.rs

1/// `z[i]`: `a[i..]` と `a` との最長共通接頭辞の長さ、を返します。
2///
3/// [実装の参考資料](https://snuke.hatenablog.com/entry/2014/12/03/214243)
4///
5/// # Examples
6/// ```
7/// use z_algorithm::z_algorithm;
8///
9/// let a = "abcabc".chars().collect::<Vec<char>>();
10/// let z = z_algorithm(&a);
11/// assert_eq!(z[0], 6); // abcabc
12/// assert_eq!(z[1], 0); // bcabc
13/// assert_eq!(z[2], 0); // cabc
14/// assert_eq!(z[3], 3); // abc
15/// assert_eq!(z[4], 0); // bc
16/// assert_eq!(z[5], 0); // c
17/// ```
18#[allow(clippy::many_single_char_names)]
19pub fn z_algorithm<T>(a: &[T]) -> Vec<usize>
20where
21    T: PartialEq + std::fmt::Debug,
22{
23    let n = a.len();
24    let mut z = vec![0; n];
25    let mut i = 0;
26    for j in 1..n {
27        if j + z[j - i] < i + z[i] {
28            debug_assert_eq!(a[j..(j + z[j - i])], a[..z[j - i]]);
29            z[j] = z[j - i];
30        } else {
31            let start = j + (i + z[i]).saturating_sub(j);
32            debug_assert_eq!(a[j..start], a[..(start - j)]);
33            let end = (start..n).find(|&k| a[k - j] != a[k]).unwrap_or(n);
34            debug_assert_eq!(a[j..end], a[..(end - j)]);
35            z[j] = end - j;
36            i = j;
37        }
38    }
39    z[0] = n;
40    z
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46    use rand::prelude::*;
47    #[test]
48    fn test() {
49        let chars = ['a', 'b', 'x', 'y'];
50        let mut rng = thread_rng();
51        for _ in 0..100 {
52            let n = rng.gen_range(1, 100);
53            let s = (0..n)
54                .map(|_| *chars.choose(&mut rng).unwrap())
55                .collect::<Vec<_>>();
56            let z = z_algorithm(&s);
57            for i in 0..n {
58                assert_eq!(z[i], lcp(&s, &s[i..]));
59            }
60        }
61    }
62
63    fn lcp(a: &[char], b: &[char]) -> usize {
64        let mut i = 0;
65        while i < a.len() && i < b.len() {
66            if a[i] != b[i] {
67                break;
68            }
69            i += 1;
70        }
71        i
72    }
73}
z_algorithm/lib.rs

z_algorithm/
lib.rs
