haar_lib/algo/
zalgo.rs

1//! Z algorithm
2//!
3//! # References
4//! - <https://snuke.hatenablog.com/entry/2014/12/03/214243>
5//! - <https://qiita.com/Pro_ktmr/items/16904c9570aa0953bf05>
6//!
7//! # Problems
8//! - <https://judge.yosupo.jp/problem/zalgorithm>
9
10/// Z algorithm
11///
12/// $k \in [1, n]$について、
13/// $s = s_1 s_2 \ldots s_n$と$s$の$k$文字目以降$s_k s_{k+1} \ldots s_n$の最長共通接頭辞の長さを求める。
14pub fn zalgo<T: PartialEq>(s: &[T]) -> Vec<usize> {
15    let n = s.len();
16    let mut ret = vec![0; n];
17    let mut j = 0;
18
19    for i in 1..n {
20        if i + ret[i - j] < j + ret[j] {
21            ret[i] = ret[i - j];
22        } else {
23            let mut k = (j + ret[j]).saturating_sub(i);
24            while i + k < n && s[k] == s[i + k] {
25                k += 1;
26            }
27            ret[i] = k;
28            j = i;
29        }
30    }
31
32    ret[0] = n;
33
34    ret
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    #[test]
42    fn test() {
43        assert_eq!(zalgo("abcbcba".as_bytes()), [7, 0, 0, 0, 0, 0, 1]);
44        assert_eq!(
45            zalgo("mississippi".as_bytes()),
46            [11, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
47        );
48        assert_eq!(zalgo("ababacaca".as_bytes()), [9, 0, 3, 0, 1, 0, 1, 0, 1]);
49        assert_eq!(zalgo("aaaaa".as_bytes()), [5, 4, 3, 2, 1]);
50    }
51}