competitive_library/string/
z_algorithm.rs

1//! Z algorithm
2pub fn z_algorithm(s: &[char]) -> Vec<usize> {
3    let length = s.len();
4    let mut z_array = vec![0_usize; length];
5
6    z_array[0] = length;
7    let (mut i, mut j) = (1, 0);
8
9    while i < length {
10        while i + j < length && s[j] == s[i + j] {
11            j += 1;
12        }
13
14        z_array[i] = j;
15
16        if j == 0 {
17            i += 1;
18            continue;
19        }
20        let mut k = 1;
21        while k < j && k + z_array[k] < j {
22            z_array[i + k] = z_array[k];
23            k += 1;
24        }
25        i += k;
26        j -= k;
27    }
28    z_array
29}
30
31#[cfg(test)]
32mod tests {
33    use super::*;
34    #[test]
35    fn test_z_algorithm() {
36        let case = vec![
37            ("abcbcba", vec![7, 0, 0, 0, 0, 0, 1]),
38            ("mississippi", vec![11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),
39            ("ababacaca", vec![9, 0, 3, 0, 1, 0, 1, 0, 1]),
40            ("aaaaa", vec![5, 4, 3, 2, 1]),
41        ];
42
43        for (s, ans) in case {
44            let z = z_algorithm(&s.to_string().chars().collect::<Vec<_>>());
45            assert_eq!(z, ans);
46        }
47    }
48}