nekolib/seq/suffix_array.rs
1//! 接尾辞配列。
2
3use std::cmp::Ordering::{Equal, Greater, Less};
4use std::collections::{BTreeMap, BTreeSet};
5use std::fmt::Debug;
6use std::ops::Index;
7
8/// 接尾辞配列。
9///
10/// 文字列 $S$ の各接尾辞を辞書順でソートしたもの。
11/// より正確には、$i$ ($0\\le i\\le |S|$) を $S\[i\\dots\]$ をキーとしてソートした配列 $A$ である。
12///
13/// # Idea
14/// ## 用語の定義など
15///
16/// 文字列 $S$ について、 $S\[|S|\]$ には辞書順最小の文字が入っていると見なす。
17/// また、$S\[|S|+1\]$ には辞書順最大の文字が入っていると見なしている気がする。
18///
19/// 以下、$S\[i\\dots\]$ を接尾辞 $i$ と呼ぶ。
20/// また、例として文字列 `GTCCCGATGTCATGTCAGGA$` を考える。
21///
22/// ```text
23/// G T C C C G A T G T C A T G T C A G G A $
24/// 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
25/// ```
26///
27/// この文字列の接尾辞配列は次のようになる。
28///
29/// ```text
30/// 20 19 16 11 6 15 10 2 3 4 18 5 17 13 8 0 14 9 1 12 7
31/// ```
32///
33/// ### S-type と L-type
34///
35/// まず、接尾辞に対して _S-type_ と _L-type_ を定義する。
36/// 接尾辞 $i$ が接尾辞 $i+1$ より辞書順で小さい (smaller)
37/// とき、接尾辞 $i$ を S-type であると言う。そうでない (larger) とき、L-type
38/// であると言う。長さが異なるので、等しくなることはないことに注意せよ。
39/// なお、接尾辞 $|S|$ は S-type とする。
40/// また、簡単のため、単に「$i$ が S-type である」などと言う。
41///
42/// $i$ ($0\\le i \\lt |S|$) について次が成り立つので、$|S|$ が S-type
43/// であることと合わせて、末尾から順に線形時間で求められる。
44///
45/// - $S\[i\] \\lt S\[i+1\]$ なら、$i$ は S-type。
46/// - $S\[i\] \\gt S\[i+1\]$ なら、$i$ は L-type。
47/// - $S\[i\] = S\[i+1\]$ なら、$i$ と $i+1$ は同じ type。
48///
49/// ### LMS suffix と LMS block
50///
51/// $i-1$ が L-type であるような S-type の $i$ を、特に
52/// LMS (leftmost S-type) suffix と言う。$0$ は LMS ではないことに注意せよ。
53/// また、LMS suffix を、次の LMS suffix（存在すれば）の開始位置で区切ったものを
54/// LMS block と呼ぶ。LMS block は次のようになる。
55///
56/// ```text
57/// G T C C C G A T G T C A T G T C A G G A $
58/// 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
59/// S L *S S S L *S L *S L L *S L *S L L *S L L L *S
60/// [------------]
61/// [------]
62/// [---------]
63/// [------]
64/// [---------]
65/// [------------]
66/// []
67/// ```
68///
69/// ### バケット
70///
71/// 接尾辞配列を再掲する。ただし、各 $i$ に対して、その type と $S\[i\]$ を併記する。
72/// ただし、`*S` は LMS を表す。
73///
74/// ```text
75/// 20 19 16 11 6 15 10 2 3 4 18 5 17 13 8 0 14 9 1 12 7
76/// $ A A A A C C C C C G G G G G G T T T T T
77/// *S L *S *S *S L L *S S S L L L *S *S S L L L L L
78/// ```
79///
80/// 同じ文字で始まる接尾辞が連続していることに気づく。この一連の部分をバケットと呼ぶ。
81/// たとえば、`19`, `16`, `11`, `6` の部分を指して `A` のバケットと呼ぶことにする。
82///
83/// さらに、各文字のバケットにおいて、L-type が前半に、S-type が後半に来ていることにも気づく。
84/// これは、$S\[i\] = S\[j\] = c$ であるような接尾辞 $i$ (L-type) と $j$ (S-type)
85/// を考えたとき、接尾辞 $i$ は $c$ が 1 つ以上続いた後に $c$ より小さい文字が現れ、
86/// 接尾辞 $j$ は $c$ が 1 つ以上続いた後に $c$ より大きい文字が現れることから従う。
87///
88/// ### Induced sorting
89///
90/// この操作により、LMS block のソート順を得ることができる。この操作では、部分文字列
91/// $S\[i\\dots j\]$ (inclusive) に辞書順最大の文字を連結したものを考え[^1]、この
92/// $i$ をバケットに入れていく。$j$ については適宜説明する。ここでは、簡単のため
93/// $S\[i\\dots j\]$ に辞書順最大の文字を連結したものを指して、単に部分文字列 $S\[i\\dots j\]$
94/// と呼ぶことにする。
95///
96/// [^1]: よくある説明では suffix $S\[i\\dots\]$ だと思いますが、
97/// それだと納得のいく説明ができなかったのでこういう説明になりました。
98/// 「LMS suffix がソート順に並んでいるときに IS をすると SA
99/// が得られます」というのが示されたあと、適当な順で LMS suffix を入力して
100/// IS を実行されると、もやもやします。
101///
102/// まず、LMS suffix たちの添字 $i$ は得られているとし、それらを対応するバケットの
103/// 末尾に適当な順に入れる。このとき $j = i$ とする。
104///
105/// ```text
106/// 20 - 6 11 16 - - - - 2 - - - - 8 13 - - - - -
107/// $ A A A A C C C C C G G G G G G T T T T T
108/// *S - *S *S *S - - - - *S - - - - *S *S - - - - -
109/// ```
110///
111/// ```text
112/// 20 $X
113/// - -
114/// 6 AX
115/// 11 AX
116/// 16 AX
117/// - -
118/// - -
119/// - -
120/// - -
121/// 2 CX
122/// - -
123/// - -
124/// - -
125/// - -
126/// 8 GX
127/// 13 GX
128/// - -
129/// - -
130/// - -
131/// - -
132/// - -
133/// ```
134///
135/// バケットを前から走査し、部分文字列 $S\[i\\dots j\]$ を見つけたとき、$i-1$ が L-type
136/// であれば、文字 $S\[i-1\]$ に対応するバケットの空きの先頭に $S\[i-1\\dots j\]$ を入れる。
137///
138/// まず、この操作によりバケットに追加された部分文字列がソート済みであることを帰納的に示す。
139/// 初期状態では明らかにソート済みである。$S\[i\\dots j\]$ を見て $S\[i-1\\dots j\]$
140/// を入れるとき、同一バケット内の L-type の各 $S\[i\'-1\\dots j\'\]$ に対して
141/// $S\[i\'-1\\dots j\'\] \\lt S\[i-1\\dots j\]$ を示せば十分である（LMS については明らかなので）。
142/// これは、$S\[i\'-1\] = S\[i-1\]$ のとき $S\[i\'-1\\dots j\'\] \\lt S\[i-1\\dots j\] \\iff
143/// S\[i\'\\dots j\'\] \\lt S\[i\\dots j\]$ であることと、$i\'$ が $i$ より先に処理されているため
144/// $S\[i\'\\dots j\'\] \\lt S\[i\\dots j\]$ であることから従う。
145///
146/// 次に、L-type であるすべての $i$ がバケットに入っていることを示す。
147/// $i+1$ が LMS であれば、定義から $i$ は L-type であり、これはバケットに入れられる。
148/// L-type である各 $i$ について $S\[i\\dots j\] \\gt S\[i+1\\dots j\]$ なので、
149/// $i+1$ がバケットに入っているならば、$i$ もバケットに入れられる。
150///
151/// ```text
152/// 20 19 6 11 16 10 15 - - 2 18 5 17 - 8 13 9 14 1 7 12
153/// $ A A A A C C C C C G G G G G G T T T T T
154/// *S L *S *S *S L L - - *S L L L - *S *S L L L L L
155/// ```
156///
157/// ```text
158/// 20 $X
159/// 19 A$X
160/// 6 AX
161/// 11 AX
162/// 16 AX
163/// 10 CAX
164/// 15 CAX
165/// - -
166/// - -
167/// 2 CX
168/// 18 GA$X
169/// 5 GAX
170/// 17 GGA$X
171/// - -
172/// 8 GX
173/// 13 GX
174/// 9 TCAX
175/// 14 TCAX
176/// 1 TCX
177/// 7 TGX
178/// 12 TGX
179/// ```
180///
181/// 次に、バケットから LMS を取り除き、逆順から同様の操作を S-type について行う。
182/// すなわち、バケットを後ろから走査し、部分文字列 $S\[i\\dots j\]$ を見つけたとき、
183/// $i-1$ が S-type であれば、文字 $S\[i-1\]$ に対応するバケットの空きの末尾に
184/// $S\[i-1\\dots j\]$ を入れる。
185///
186/// L-type のときと同様の議論から、すべての S-type の $i$ がソート順で得られることがわかる。
187/// S-type と L-type の性質から、すべての $i$ についてソート順で得られたことになる。
188///
189/// ```text
190/// 20 19 16 6 11 10 15 2 3 4 18 5 17 8 13 0 9 14 1 7 12
191/// $ A A A A C C C C C G G G G G G T T T T T
192/// *S L *S *S *S L L *S S S L L L *S *S S L L L L L
193/// ```
194///
195/// ```text
196/// 20 $X
197/// 19 A$X
198/// 16 AGGA$X
199/// 6 ATGX
200/// 11 ATGX
201/// 10 CAX
202/// 15 CAX
203/// 2 CCCGAX
204/// 3 CCGAX
205/// 4 CGAX
206/// 18 GA$X
207/// 5 GAX
208/// 17 GGA$X
209/// 8 GTCAX
210/// 13 GTCAX
211/// 0 GTCX
212/// 9 TCAX
213/// 14 TCAX
214/// 1 TCX
215/// 7 TGX
216/// 12 TGX
217/// ```
218///
219/// 各 LMS $i$ に対応する部分文字列が LMS block と等しくなることは簡単に示せて、
220/// 以上より LMS block のソート順が得られることがわかった。なお、初めに LMS
221/// をバケットに入れるときにソート順に入れていれば、$j = |S|$
222/// としても同様の議論ができ、接尾辞配列が得られる。
223///
224/// ## アルゴリズムの概要
225///
226/// まず、$S$ を末尾から走査し、type を求める。
227///
228/// ```text
229/// G T C C C G A T G T C A T G T C A G G A $
230/// 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
231/// S L *S S S L *S L *S L L *S L *S L L *S L L L *S
232/// ```
233///
234/// 次に、LMS を出現順にバケットの末尾に置いて induced sorting を行う。
235///
236/// ```text
237/// 20 - 6 11 16 - - - - 2 - - - - 8 13 - - - - -
238/// $ A A A A C C C C C G G G G G G T T T T T
239/// *S - *S *S *S - - - - *S - - - - *S *S - - - - -
240/// ```
241///
242/// ```text
243/// 20 19 6 11 16 10 15 - - 2 18 5 17 - 8 13 9 14 1 7 12
244/// $ A A A A C C C C C G G G G G G T T T T T
245/// *S L *S *S *S L L - - *S L L L - *S *S L L L L L
246/// ```
247///
248/// ```text
249/// 20 19 16 6 11 10 15 2 3 4 18 5 17 8 13 0 9 14 1 7 12
250/// $ A A A A C C C C C G G G G G G T T T T T
251/// *S L *S *S *S L L *S S S L L L *S *S S L L L L L
252/// ```
253///
254/// これにより、LMS block の正しいソート順が得られる。
255/// `*S` の添字を得られた順に並べると `20 16 6 11 2 8 13` となり、対応する
256/// LMS block を順に並べると次のようになる[^2]。
257///
258/// [^2]: `$` で始まる `<pre>` がシェルスクリプトではない例。
259///
260/// ```text
261/// $ AGGA$ ATG ATG CCCGA GTCA GTCA
262/// ```
263///
264/// これはソート順で得られているので、隣同士の要素を比較するだけで、各要素の順序づけがわかる。
265/// 等しい要素があることに注意せよ。
266///
267/// ```text
268/// G T C C C G A T G T C A T G T C A G G A $
269/// 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
270/// S L *S S S L *S L *S L L *S L *S L L *S L L L *S
271/// [----------][----][-------][----][-------][----------][]
272/// 3 2 4 2 4 1 0
273/// ```
274///
275/// この順序づけによってできる列 (_reduced string_) `3 2 4 2 4 1 0` の接尾辞配列は
276/// 再帰的に SA-IS によって求めることができ、`6 5 3 1 0 4 2` となる。
277///
278/// LMS block たちをそのまま連結させると境界部分の文字が重複するのが気になる。
279/// reduced string 中にも暗黙にそれらの重複があるが、それが許されることに触れておく。
280/// 等しいペアについては、重複の 1 文字目で等しければ 2 文字目でも等しいので、
281/// 辞書順比較には影響しない。等しくないペアについて、`$` 以外の LMS block の形は
282/// `SS..SLL..LS` であり、最悪時でも片方に `LS` が現れた時点で大小関係が決まるので、
283/// 重複部分まで走査されることがない。
284///
285/// これを元の文字列に対応させることで、LMS suffix のソート順が `20 16 11 6 2 13 8`
286/// であるとわかり、これを用いて再度 induced sorting を行う。
287///
288/// ```text
289/// 20 - 16 11 6 - - - - 2 - - - - 13 8 - - - - -
290/// $ A A A A C C C C C G G G G G G T T T T T
291/// *S - *S *S *S - - - - *S - - - - *S *S - - - - -
292/// ```
293///
294/// ```text
295/// 20 19 16 11 6 15 10 - - 2 18 5 17 - 13 8 14 9 1 12 7
296/// $ A A A A C C C C C G G G G G G T T T T T
297/// *S L *S *S *S L L - - *S L L L - *S *S L L L L L
298/// ```
299///
300/// ```text
301/// 20 19 16 11 6 15 10 2 3 4 18 5 17 13 8 0 14 9 1 12 7
302/// $ A A A A C C C C C G G G G G G T T T T T
303/// *S L *S *S *S L L *S S S L L L *S *S S L L L L L
304/// ```
305///
306/// これにより、所望の接尾辞配列が得られる。
307///
308/// 再帰の base case は、列の各要素が distinct な場合で、逆順列が接尾辞配列となる。
309/// 連続して LMS suffix が現れないことと、先頭位置は LMS にならないことから、
310/// reduced string は元の文字列の半分以下の長さになることがわかる。
311/// よって、長さ $n$ での計算量を $T(n)$ とおくと、$T(1) \\in O(1)$ より次のようになる。
312/// $$ T(n) \\le T(n/2) + O(n) \\in O(n). $$
313///
314/// ## パターン検索
315///
316/// パターン文字列 $T$ の $S$ における出現位置を検索する。
317///
318/// 範囲 $\[0\\dots |S|\]$ を接尾辞配列 $A$ 上で二分探索する。
319/// これは、$S\[A\[i\]\\dots\]$ が $i$ に対して単調増加であることから可能。
320/// 境界の両端を求めることで全列挙したり、片側のみを求めてパターンの有無のみを調べたりできる。
321///
322/// # Complexity
323/// 入力中の文字の種類を $\\sigma$、文字列長を $n$ とする。
324/// SA-IS を用いて構築するため、前処理は $O(\\sigma\\log(\\sigma)+n)$ 時間。
325///
326/// ↑ `todo!()` `String` なら $O(\\sigma+n)$、`[T]` なら mapping に時間がかかって $O(n\\log(\\sigma))$
327/// っぽそう。値の範囲が `0..n` に収まるような `[usize]` に対しては $O(n)$ でできるから、追記かも。
328///
329/// 検索は、パターン長を $m$ として $O(m\\log(n))$ 時間。
330///
331/// # References
332/// - Nong, Ge, Sen Zhang, and Wai Hong Chan. "Two efficient algorithms for linear time suffix array construction." _IEEE Transactions on Computers_ 60, no. 10 (2010): 1471--1484.
333/// - Ko, Pang, and Srinivas Aluru. "Space efficient linear time construction of suffix arrays." In _Annual Symposium on Combinatorial Pattern Matching_, pp. 200--210. Springer, Berlin, Heidelberg, 2003.
334///
335/// ## See also
336/// [CS166](http://web.stanford.edu/class/archive/cs/cs166/cs166.1206/lectures/04/Slides04.pdf)。
337/// 差分スライドの関係でページ数がめちゃくちゃ多くて重いので注意。軽いのは
338/// [こっち](http://web.stanford.edu/class/archive/cs/cs166/cs166.1206/lectures/04/Small04.pdf)。
339#[derive(Clone, Debug, Eq, PartialEq)]
340pub struct SuffixArray<T: Ord> {
341 buf: Vec<T>,
342 sa: Vec<usize>,
343}
344
345impl<T: Ord> From<Vec<T>> for SuffixArray<T> {
346 fn from(buf: Vec<T>) -> Self {
347 let buf_usize = hash(&buf);
348 let sa = sa_is(&buf_usize);
349 Self { buf, sa }
350 }
351}
352
353impl From<String> for SuffixArray<char> {
354 fn from(buf: String) -> Self {
355 let buf: Vec<_> = buf.as_str().chars().collect();
356 let buf_usize = hash_chars(&buf);
357 let sa = sa_is(&buf_usize);
358 Self { buf, sa }
359 }
360}
361
362impl SuffixArray<u8> {
363 pub fn from_bytes(buf: Vec<u8>) -> Self {
364 let buf_usize = hash_bytes(&buf);
365 let sa = sa_is(&buf_usize);
366 Self { buf, sa }
367 }
368}
369
370impl SuffixArray<usize> {
371 pub fn from_hashed(buf: Vec<usize>) -> Self {
372 assert!(Self::is_hashed(&buf));
373 let buf_usize: Vec<_> =
374 buf.iter().map(|&x| x + 1).chain(Some(0)).collect();
375
376 let sa = sa_is(&buf_usize);
377 Self { buf, sa }
378 }
379
380 fn is_hashed(buf: &[usize]) -> bool {
381 let mut count = vec![0; buf.len()];
382 for &x in buf {
383 count[x] += 1;
384 }
385 (0..buf.len())
386 .find(|&i| count[i] == 0)
387 .map(|i| (i..count.len()).all(|i| count[i] == 0))
388 .unwrap_or(true)
389 }
390}
391
392/// 座標圧縮をする。
393///
394/// `buf` の末尾に辞書順最小の文字 `$` を付加した列を、座標圧縮して返す。
395fn hash<T: Ord>(buf: &[T]) -> Vec<usize> {
396 let enc: BTreeSet<_> = buf.iter().collect();
397 let enc: BTreeMap<_, _> =
398 enc.into_iter().enumerate().map(|(i, x)| (x, i)).collect();
399 buf.iter()
400 .map(|x| enc[x] + 1)
401 .chain(std::iter::once(0)) // for '$'
402 .collect()
403}
404
405/// 座標圧縮をする。
406///
407/// `buf` の末尾に辞書順最小の文字 `$` を付加した列を、座標圧縮して返す。`char`
408/// の列を受け取り、バケットソートの要領で行う。
409fn hash_chars(buf: &[char]) -> Vec<usize> {
410 let max = match buf.iter().max() {
411 Some(&c) => c as usize,
412 None => return vec![0], // "$"
413 };
414 let enc = {
415 let mut enc = vec![0; max + 1];
416 for &c in buf {
417 enc[c as usize] = 1;
418 }
419 for i in 1..=max {
420 enc[i] += enc[i - 1];
421 }
422 enc
423 };
424 buf.iter().map(|&x| enc[x as usize])
425 .chain(std::iter::once(0)) // for '$'
426 .collect()
427}
428
429/// 座標圧縮をする。
430///
431/// `buf` の末尾に辞書順最小の文字 `$` を付加した列を、座標圧縮して返す。`u8`
432/// の列を受け取り、バケットソートの要領で行う。
433fn hash_bytes(buf: &[u8]) -> Vec<usize> {
434 let enc = {
435 let mut enc = vec![0; 256];
436 for &b in buf {
437 enc[b as usize] = 1;
438 }
439 for i in 1..=255 {
440 enc[i] += enc[i - 1];
441 }
442 enc
443 };
444 buf.iter().map(|&x| enc[x as usize] )
445 .chain(std::iter::once(0)) // for '$'
446 .collect()
447}
448
449#[derive(Clone, Copy, Debug, Eq, PartialEq)]
450enum LsType {
451 LType,
452 SType(bool), // true iff leftmost S-type
453}
454use LsType::{LType, SType};
455
456/// 出現回数を求める。
457///
458/// `res[x]` が `buf` における `x` の出現回数となる配列 `res` を返す。
459///
460/// # Requirements
461/// `buf` の要素は `0..buf.len()` に含まれる。
462fn count_freq(buf: &[usize]) -> Vec<usize> {
463 let mut res = vec![0; buf.len()];
464 buf.iter().for_each(|&x| res[x] += 1);
465 res
466}
467
468/// 逆順列を返す。
469///
470/// `res[x]` が `buf` における `x` の出現位置となる配列 `res` を返す。
471///
472/// # Requirements
473/// `buf` の要素は `0..buf.len()` に含まれ、かつ distinct である。
474fn inv_perm(buf: &[usize]) -> Vec<usize> {
475 let mut res = vec![0; buf.len()];
476 buf.iter().enumerate().for_each(|(i, &x)| res[x] = i);
477 res
478}
479
480/// LS type を求める。
481///
482/// `res[i]` が `buf[i]` の LS type である配列 `res` を返す。
483fn ls_classify(buf: &[usize]) -> Vec<LsType> {
484 let mut res = vec![SType(false); buf.len()];
485 for i in (0..buf.len() - 1).rev() {
486 res[i] = match buf[i].cmp(&buf[i + 1]) {
487 Less => SType(false),
488 Equal => res[i + 1],
489 Greater => LType,
490 };
491 }
492 for i in 1..buf.len() {
493 if let (LType, SType(_)) = (res[i - 1], res[i]) {
494 res[i] = SType(true);
495 }
496 }
497 res
498}
499
500/// 各バケットの開始位置を求める。
501///
502/// # Input
503/// `count[i]` が `i` 番目のバケットのサイズである配列 `count`。
504fn bucket_head(count: &[usize]) -> Vec<usize> {
505 let n = count.len();
506 let mut head: Vec<_> =
507 std::iter::once(&0).chain(&count[..n - 1]).cloned().collect();
508 for i in 1..n {
509 head[i] += head[i - 1];
510 }
511 head
512}
513
514/// 各バケットの終端位置を求める。
515///
516/// # Input
517/// `count[i]` が `i` 番目のバケットのサイズである配列 `count`。
518fn bucket_tail(count: &[usize]) -> Vec<usize> {
519 let mut tail = count.to_vec();
520 for i in 1..count.len() {
521 tail[i] += tail[i - 1];
522 }
523 tail
524}
525
526/// すでに判明している SA と LS type から、SA の残りの要素を求める。
527fn induce(
528 buf: &[usize],
529 sa: &mut [Option<usize>],
530 count: &[usize],
531 ls: &[LsType],
532) {
533 let mut head = bucket_head(count);
534 for i in 0..sa.len() {
535 if let Some(j) = sa[i] {
536 if j > 0 && ls[j - 1] == LType {
537 sa[head[buf[j - 1]]] = Some(j - 1);
538 head[buf[j - 1]] += 1;
539 }
540 }
541 }
542 let mut tail = bucket_tail(count);
543 for i in (1..count.len()).rev() {
544 if let Some(j) = sa[i] {
545 if j > 0 && ls[j - 1] != LType {
546 tail[buf[j - 1]] -= 1;
547 sa[tail[buf[j - 1]]] = Some(j - 1);
548 }
549 }
550 }
551}
552
553/// 各 LMS block を文字とする文字列を作る。
554///
555/// 新たな文字も辞書順に番号づけられる。
556///
557/// # Examples
558/// `[CCCG][AT][GTC][AT][GTC][AGGA][$]` → `3242410`
559fn reduce(buf: &[usize], lms: &[usize], ls: &[LsType]) -> Vec<usize> {
560 if lms.len() <= 1 {
561 return vec![0; lms.len()];
562 }
563
564 let e = |(i0, i1)| {
565 if (ls[i0], ls[i1]) == (SType(true), SType(true)) {
566 Some(true)
567 } else if ls[i0] != ls[i1] || buf[i0] != buf[i1] {
568 Some(false)
569 } else {
570 None
571 }
572 };
573
574 let mut map = vec![0; buf.len()]; // map[lms[0]] = 0
575 map[lms[1]] = 1;
576 let mut x = 1;
577 for i in 2..lms.len() {
578 let equiv = buf[lms[i]] == buf[lms[i - 1]]
579 && (lms[i] + 1..).zip(lms[i - 1] + 1..).find_map(e).unwrap();
580 if !equiv {
581 x += 1;
582 }
583 map[lms[i]] = x;
584 }
585
586 (0..buf.len())
587 .filter_map(|i| match ls[i] {
588 SType(true) => Some(map[i]),
589 _ => None,
590 })
591 .collect()
592}
593
594/// SA-IS により接尾辞配列を求める。
595fn sa_is(buf: &[usize]) -> Vec<usize> {
596 let len = buf.len();
597 let count = count_freq(buf);
598 if count.iter().all(|&x| x == 1) {
599 return inv_perm(buf);
600 }
601
602 let ls = ls_classify(buf);
603 let mut sa = vec![None; len];
604 let mut tail = bucket_tail(&count);
605 for i in (1..len).rev().filter(|&i| ls[i] == SType(true)) {
606 tail[buf[i]] -= 1;
607 sa[tail[buf[i]]] = Some(i);
608 }
609
610 induce(buf, &mut sa, &count, &ls);
611
612 let lms: Vec<_> = sa
613 .into_iter()
614 .map(std::option::Option::unwrap)
615 .filter(|&i| ls[i] == SType(true))
616 .collect(); // in lexicographic order
617 let rs_sa = sa_is(&reduce(buf, &lms, &ls));
618
619 // in appearing order
620 let lms: Vec<_> = (0..len).filter(|&i| ls[i] == SType(true)).collect();
621
622 let mut tail = bucket_tail(&count);
623 let mut sa = vec![None; len];
624 for i in rs_sa.into_iter().rev() {
625 let j = lms[i];
626 tail[buf[j]] -= 1;
627 sa[tail[buf[j]]] = Some(j);
628 }
629 induce(buf, &mut sa, &count, &ls);
630
631 sa.into_iter().map(std::option::Option::unwrap).collect()
632}
633
634impl<T: Ord> SuffixArray<T> {
635 /// パターン検索を行う。
636 ///
637 /// # Complexity
638 ///
639 /// $O(|T|\\log(|S|))$ 時間。
640 ///
641 /// # Examples
642 ///
643 /// ```
644 /// use nekolib::seq::SuffixArray;
645 ///
646 /// let s: Vec<_> = "abracadabra".chars().collect();
647 /// let sa: SuffixArray<_> = s.into();
648 ///
649 /// assert_eq!(sa.search(&['a']).collect::<Vec<_>>(), vec![10, 7, 0, 3, 5]);
650 /// assert_eq!(
651 /// sa.search(&"abra".chars().collect::<Vec<_>>()).nth(1),
652 /// Some(0)
653 /// );
654 /// assert_eq!(sa.search(&['a', 'e']).next(), None);
655 /// ```
656 pub fn search(&self, pat: &[T]) -> impl Iterator<Item = usize> + '_ {
657 let lo = {
658 let mut lt = 1_usize.wrapping_neg();
659 let mut ge = self.sa.len();
660 while ge.wrapping_sub(lt) > 1 {
661 let mid = lt.wrapping_add(ge.wrapping_sub(lt) / 2);
662 let pos = self.sa[mid];
663 match self.buf[pos..].cmp(pat) {
664 Less => lt = mid,
665 _ => ge = mid,
666 }
667 }
668 ge
669 };
670 if lo >= self.sa.len() {
671 return self.sa[lo..lo].iter().cloned();
672 }
673 let hi = {
674 let mut le = lo.wrapping_sub(1);
675 let mut gt = self.sa.len();
676 while gt.wrapping_sub(le) > 1 {
677 let mid = le.wrapping_add(gt.wrapping_sub(le) / 2);
678 let pos = self.sa[mid];
679 let len = pat.len().min(self.buf[pos..].len());
680 match self.buf[pos..pos + len].cmp(pat) {
681 Greater => gt = mid,
682 _ => le = mid,
683 }
684 }
685 gt
686 };
687 self.sa[lo..hi].iter().cloned()
688 }
689
690 /// 高さ配列を返す。
691 ///
692 /// # Examples
693 /// ```
694 /// use nekolib::seq::SuffixArray;
695 ///
696 /// let s: Vec<_> = "abracadabra".chars().collect();
697 /// let sa: SuffixArray<_> = s.into();
698 /// assert_eq!(sa.lcpa(), [0, 0, 1, 4, 1, 1, 0, 3, 0, 0, 0, 2]);
699 /// ```
700 pub fn lcpa(&self) -> Vec<usize> {
701 let n = self.buf.len();
702 let mut rank = vec![0; n + 1];
703 for i in 0..=n {
704 rank[self.sa[i]] = i;
705 }
706 let mut h = 0;
707 let mut res = vec![0; n + 1];
708 for i in 0..n {
709 let j = self.sa[rank[i] - 1];
710 if h > 0 {
711 h -= 1;
712 }
713 while j + h < n && i + h < n {
714 if self.buf[j + h] != self.buf[i + h] {
715 break;
716 }
717 h += 1;
718 }
719 res[rank[i]] = h;
720 }
721 res
722 }
723
724 /// 自身を消費し、内部表現を返す。
725 ///
726 /// # Examples
727 ///
728 /// ```
729 /// use nekolib::seq::SuffixArray;
730 ///
731 /// let s: Vec<_> = "abracadabra".chars().collect();
732 /// let sa: SuffixArray<_> = s.into();
733 /// let sa = sa.into_inner();
734 /// assert_eq!(sa, vec![11, 10, 7, 0, 3, 5, 8, 1, 4, 6, 9, 2]);
735 /// ```
736 pub fn into_inner(self) -> Vec<usize> { self.sa }
737}
738
739impl SuffixArray<char> {
740 /// パターン文字列検索を行う。
741 ///
742 /// # Examples
743 /// ```
744 /// use nekolib::seq::SuffixArray;
745 ///
746 /// let sa: SuffixArray<_> = "abracadabra".to_string().into();
747 /// let occ: Vec<_> = sa.search_str("ab").collect();
748 /// assert_eq!(occ, vec![7, 0]);
749 /// let occ: Vec<_> = sa.search_str("a").collect();
750 /// assert_eq!(occ, vec![10, 7, 0, 3, 5]);
751 /// assert_eq!(sa.search_str("e").next(), None);
752 /// ```
753 pub fn search_str(&self, pat: &str) -> impl Iterator<Item = usize> + '_ {
754 let pat: Vec<_> = pat.chars().collect();
755 self.search(&pat)
756 }
757}
758
759impl<T: Ord> Index<usize> for SuffixArray<T> {
760 type Output = usize;
761 fn index(&self, i: usize) -> &usize { &self.sa[i] }
762}
763
764impl<T: Ord> From<SuffixArray<T>> for Vec<usize> {
765 fn from(sa: SuffixArray<T>) -> Self { sa.sa }
766}
767
768#[test]
769fn test_simple() {
770 let buf = "abracadabra".to_string();
771 let sa: SuffixArray<_> = buf.into();
772 let sa: Vec<_> = sa.into();
773 assert_eq!(sa, [11, 10, 7, 0, 3, 5, 8, 1, 4, 6, 9, 2])
774}
775
776#[test]
777fn test_empty_text() {
778 let sa: SuffixArray<_> = "".to_string().into();
779 let occ: Vec<_> = sa.search_str("").collect();
780 assert_eq!(occ, [0]);
781 assert_eq!(sa.search_str("x").next(), None);
782}
783
784#[test]
785fn test_empty_pat() {
786 let sa: SuffixArray<_> = "empty".to_string().into();
787 let occ: Vec<_> = sa.search_str("").collect();
788 assert_eq!(occ, [5, 0, 1, 2, 3, 4]);
789}
790
791#[test]
792fn test_naive() {
793 let n = 1000;
794 let f = |x: &i32| Some((x * 29 + 71) % 143);
795 let buf: Vec<_> = std::iter::successors(Some(2_i32), f).take(n).collect();
796 let naive_sa = {
797 let mut sa: Vec<_> = (0..=n).collect();
798 sa.sort_unstable_by_key(|&i| &buf[i..]);
799 sa
800 };
801 let sa: SuffixArray<_> = buf.into();
802 let sa: Vec<_> = sa.into();
803 assert_eq!(sa, naive_sa);
804}