Skip to main content

nekolib/seq/
z_algo.rs

1//! Z algorithm。
2
3use std::fmt::Debug;
4use std::ops::Range;
5
6/// Z algorithm。
7///
8/// 文字列 $S$ について、$Z\[i\]$ ($0\\le i < |S|$) が
9/// $S$ と $S\[i\\dots\]$ の最長共通接頭辞の長さであるような配列 $Z$ を構築する。
10///
11/// # Implementation notes
12/// テキスト `T` 中のパターン `P` を探すとき、`T` と `P` に含まれない文字 `'$'`
13/// を用いて作った文字列 `P + '$' + T` の Z value を計算することで求められる。
14///
15/// 一方、この方法においては `'$'` を適切に探す必要がある。
16/// caller 側が探すのは面倒・バグの元であり、callee 側が探すのはコストがかかる。
17/// あるいは、そういう元が存在しない場合もなくはない。
18///
19/// さて、`P + '$' + T` を読み込んだ場合の挙動は `'$'` を仮定しなくても模倣できる。
20/// よって、多少の実装量には目をつぶりつつそういう実装にした。
21/// テキストがパターン構築時には与えられないような場合や、
22/// 複数の処理がオンラインに与えられる場合にも対応しやすいと思う。
23///
24/// # Suggestions
25/// [`std::str::matches`](https://doc.rust-lang.org/std/primitive.str.html#method.matches)
26/// の実装を参考にするなどして、`fn occurrences(&self)` を持つような trait
27/// を作るとよさそう。テストを KMP に対して再利用できるので。
28///
29/// # Complexity
30/// 構築は $O(|S|)$ 時間。検索は、テキスト $T$ に対して $O(|T|)$ 時間。
31#[derive(Clone, Debug, Eq, PartialEq)]
32pub struct ZSearcher<T: Eq> {
33    pat: Vec<T>,
34    z: Vec<usize>,
35}
36
37impl<T: Clone + Eq> From<Vec<T>> for ZSearcher<T> {
38    fn from(pat: Vec<T>) -> Self {
39        let len = pat.len();
40        let mut z = vec![len; len];
41        let mut i = 1;
42        let mut j = 0;
43        while i < len {
44            while i + j < len && pat[j] == pat[i + j] {
45                j += 1;
46            }
47            z[i] = j;
48            if j == 0 {
49                i += 1;
50                continue;
51            }
52            let mut k = 1;
53            while i + k < len && k + z[k] < j {
54                z[i + k] = z[k];
55                k += 1;
56            }
57            i += k;
58            j -= k;
59        }
60        Self { pat, z }
61    }
62}
63
64impl<T: Eq> ZSearcher<T> {
65    pub fn occurrences<'a>(&'a self, s: &'a [T]) -> Occurrences<'a, T> {
66        Occurrences { text_index: 0, match_len: 0, z: &self, text: s }
67    }
68
69    pub fn z(&self, i: usize) -> usize { self.z[i] }
70}
71
72pub struct Occurrences<'a, T: Eq> {
73    text_index: usize,
74    match_len: usize,
75    z: &'a ZSearcher<T>,
76    text: &'a [T],
77}
78
79impl<T: Eq> Iterator for Occurrences<'_, T> {
80    type Item = Range<usize>;
81    fn next(&mut self) -> Option<Self::Item> {
82        let text = self.text;
83        let pat = &self.z.pat;
84        let z = &self.z.z;
85
86        if pat.is_empty() {
87            return if self.text_index < text.len() {
88                let i = self.text_index;
89                self.text_index += 1;
90                Some(i..i)
91            } else {
92                None
93            };
94        }
95
96        let mut i = self.text_index;
97        let mut j = self.match_len;
98        while i < text.len() {
99            while i < text.len() && j < pat.len() && pat[j] == text[i] {
100                i += 1;
101                j += 1;
102            }
103            if j == 0 {
104                i += 1;
105                continue;
106            }
107            let mut k = 1;
108            while k < pat.len() && k + z[k] < j {
109                k += 1;
110            }
111            if j == pat.len() {
112                self.text_index = i;
113                self.match_len = j - k;
114                return Some(i - j..i);
115            }
116            j -= k;
117        }
118        self.text_index = text.len();
119        None
120    }
121}