competitive_programming_rs/string/
z_algorithm.rs

1pub mod z_algorithm {
2
3    pub fn calc_z_array<T: PartialEq>(s: &[T]) -> Vec<usize> {
4        let n = s.len();
5        let mut z_array = vec![0; n];
6
7        // [l, r) is a window which matches with prefix of s
8        let mut l = 0;
9        let mut r = 1;
10        for i in 1..n {
11            if i >= r {
12                l = i;
13                r = i + 1;
14                while r <= n && s[r - 1 - l] == s[r - 1] {
15                    r += 1;
16                }
17                z_array[i] = r - l - 1;
18                r -= 1;
19            } else if z_array[i - l] < r - i {
20                z_array[i] = z_array[i - l];
21            } else {
22                l = i;
23                while r <= n && s[r - 1 - l] == s[r - 1] {
24                    r += 1;
25                }
26                z_array[i] = r - l - 1;
27                r -= 1;
28            }
29        }
30        z_array[0] = n;
31        z_array
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use super::*;
38    use rand::distributions::Uniform;
39    use rand::Rng;
40
41    #[test]
42    fn test_z() {
43        let n = 1000;
44        let mut rng = rand::thread_rng();
45
46        for _ in 0..100 {
47            let mut s = String::new();
48            for _ in 0..n {
49                let c = (rng.sample(Uniform::from(0..26)) as u8 + 'a' as u8) as char;
50                s.push(c);
51            }
52
53            let t = String::new() + s.as_str() + s.as_str();
54
55            let z_array = z_algorithm::calc_z_array(&t.as_bytes());
56            let n = t.len();
57            for i in 0..n {
58                let l = z_array[i];
59                assert_eq!(&t[0..l], &t[i..(i + l)]);
60                assert!(i + l >= t.len() || t[l..(l + 1)] != t[(i + l)..(i + l + 1)]);
61            }
62        }
63    }
64}
competitive_programming_rs/string/z_algorithm.rs

competitive_programming_rs/string/
z_algorithm.rs
