Skip to main content

nekolib/seq/
kmp.rs

1//! KMP 法。
2
3use super::super::traits::push_pop;
4
5use std::fmt::Debug;
6use std::ops::Range;
7
8use push_pop::{PopBack, PushBack};
9
10/// KMP 法 (Knuth–Morris–Pratt algorithm)。
11///
12/// 文字列 $S$ を入力とする。
13/// 各 $S\[\\dots i\]$ ($0\\le i\\le |S|$) の最長 border と最長 tagged border
14/// の長さを求める。該当する border が存在しない要素は $-1$ とする。
15///
16/// 文字列 $S\[\\dots i\]$ の _border_ とは、$S\[\\dots i\]$ の真部分文字列であり、$S\[\\dots i\]$
17/// の接頭辞かつ接尾辞であるような文字列である。
18///
19/// 文字列 $S\[\\dots i\]$ の border $S\[\\dots j\]$ ($0\\le j < i$) が $S\[j\] \\neq S\[i\]$
20/// を満たすとき、$S\[\\dots j\]$ は $S\[\\dots i\]$ の _tagged border_ (_strict border_,
21/// _strong border_) であると言う。ただし、$S\[|S|\]$ は $S$ 中に含まれない文字として定義する。
22///
23/// ```text
24///     0 1 2 3 4 5 6 7 8   9 ...
25///   +-------------------+-------+
26/// S | a a b a c a a b a | c ... |
27///   +-------------------+-------+
28/// ```
29///
30/// この例において、$S\[\\dots 4\] = \\mathtt{aaba}$ は $S\[\\dots 9\] = \\mathtt{aabacaaba}$
31/// の border だが tagged border ではない ($S\[4\] = S\[9\] = \\mathtt{c}$)。一方、
32/// $S\[\\dots 2\] = \\mathtt{aa}$ は tagged border である
33/// ($S\[2\] = \\mathtt{b} \\neq S\[9\] = \\mathtt{c}$)。
34///
35/// この tagged border を用いることで、パターン検索を高速に行う。
36///
37/// # Idea
38/// tagged border を $O(|S|)$ 時間で求める方法と、パターン検索を $O(|T|)$
39/// 時間で行う方法について書く。
40///
41/// # Implementation notes
42/// パターンが静的な場合であれば最長 border を求める必要はない。
43/// パターンの末尾に対する push のみを考慮する場合も同様のはず。
44/// ここではパターンの末尾に対する pop を行える実装になっており、その更新の際には最長
45/// border の長さが必要になるはず。
46///
47/// # Complexity
48/// 構築は $O(|S|)$ 時間。パターン末尾における更新は $O(\\log(|S|))$ 時間。
49/// ただし、末尾への追加はこれに worst $O(|S|)$ (amortized $O(1)$) 時間が加わる。
50///
51/// # References
52/// - Knuth, D. E., Morris, Jr, J. H., & Pratt, V. R. (1977). Fast pattern matching in strings. _SIAM journal on computing, 6_(2), 323-350.
53///
54/// # See also
55/// - <https://snuke.hatenablog.com/entry/2014/12/01/235807>
56/// - <https://snuke.hatenablog.com/entry/2017/07/18/101026>
57/// - <https://potetisensei.hatenablog.com/entry/2017/07/10/174908>
58/// - <http://www-igm.univ-mlv.fr/~lecroq/string/node8.html>
59#[derive(Clone, Debug, Eq, PartialEq)]
60pub struct KmpSearcher<T: Eq> {
61    pat: Vec<T>,
62    fail: Vec<usize>,
63    fail1: Vec<usize>,
64}
65
66impl<T: Eq> From<Vec<T>> for KmpSearcher<T> {
67    fn from(pat: Vec<T>) -> Self {
68        let len = pat.len();
69        let fail1 = vec![1_usize.wrapping_neg(); len + 1];
70        let fail = fail1.clone();
71        let mut self_ = Self { pat, fail1, fail };
72        for i in 0..self_.pat.len() {
73            let (f, f1) = self_.calc_fail(i);
74            self_.fail1[i + 1] = f1;
75            self_.fail[i + 1] = f;
76        }
77        self_
78    }
79}
80
81impl<T: Eq> KmpSearcher<T> {
82    fn calc_fail(&self, i: usize) -> (usize, usize) {
83        let pat_i = &self.pat[i];
84        let mut j = self.fail1[i];
85        while j < self.pat.len() && pat_i != &self.pat[j] {
86            j = self.fail[j];
87        }
88        j = j.wrapping_add(1);
89        match self.pat.get(i + 1) {
90            Some(pat_ni) if pat_ni == &self.pat[j] => (self.fail[j], j),
91            _ => (j, j),
92        }
93    }
94
95    pub fn occurrences<'a>(&'a self, s: &'a [T]) -> Occurrences<'a, T> {
96        Occurrences {
97            text_index: 0,
98            pat_index: 0,
99            kmp: &self,
100            text: s,
101        }
102    }
103}
104
105pub struct Occurrences<'a, T: Eq> {
106    text_index: usize,
107    pat_index: usize,
108    kmp: &'a KmpSearcher<T>,
109    text: &'a [T],
110}
111
112impl<T: Eq> Iterator for Occurrences<'_, T> {
113    type Item = Range<usize>;
114    fn next(&mut self) -> Option<Self::Item> {
115        let text = self.text;
116        let pat = &self.kmp.pat;
117
118        if pat.is_empty() {
119            return if self.text_index < text.len() {
120                let i = self.text_index;
121                self.text_index += 1;
122                Some(i..i)
123            } else {
124                None
125            };
126        }
127
128        let mut j = self.pat_index;
129        for (i, c) in text[self.text_index..].iter().enumerate() {
130            let i = i + self.text_index;
131            while j < pat.len() && &pat[j] != c {
132                j = self.kmp.fail[j];
133            }
134            j = j.wrapping_add(1);
135            if j == pat.len() {
136                let e = i + 1;
137                let res = e - pat.len()..e;
138                self.text_index = e;
139                self.pat_index = self.kmp.fail[j];
140                return Some(res);
141            }
142        }
143        self.text_index = text.len();
144        None
145    }
146}
147
148impl<T: Eq> PushBack for KmpSearcher<T> {
149    type Input = T;
150    fn push_back(&mut self, x: T) {
151        let len = self.pat.len();
152        if len > 0 {
153            let j = self.fail1[len];
154            if &self.pat[j] == &x {
155                let tmp = self.fail[j];
156                self.fail[len] = tmp;
157            }
158        }
159        self.pat.push(x);
160        let (f, f1) = self.calc_fail(len);
161        self.fail1.push(f1);
162        self.fail.push(f);
163    }
164}
165
166impl<T: Eq> PopBack for KmpSearcher<T> {
167    type Output = usize;
168    fn pop_back(&mut self) -> Option<usize> {
169        if self.pat.is_empty() {
170            None
171        } else {
172            self.pat.pop();
173            self.fail1.pop();
174            let res = self.fail.pop();
175            *self.fail.last_mut().unwrap() = *self.fail1.last().unwrap();
176            res
177        }
178    }
179}