competitive/string/
suffix_automaton.rs

1use std::collections::HashMap;
2
3#[derive(Debug)]
4struct State {
5    next: HashMap<usize, usize>,
6    link: usize,
7    len: usize,
8}
9
10#[derive(Debug)]
11pub struct SuffixAutomaton {
12    states: Vec<State>,
13    last: usize,
14}
15
16impl Default for SuffixAutomaton {
17    fn default() -> Self {
18        Self {
19            states: vec![State {
20                next: HashMap::new(),
21                link: !0,
22                len: 0,
23            }],
24            last: 0,
25        }
26    }
27}
28
29impl SuffixAutomaton {
30    pub fn new() -> Self {
31        Self::default()
32    }
33
34    pub fn state_size(&self) -> usize {
35        self.states.len()
36    }
37
38    pub fn transitions(&self, state: usize) -> &HashMap<usize, usize> {
39        &self.states[state].next
40    }
41
42    pub fn length(&self, state: usize) -> usize {
43        self.states[state].len
44    }
45
46    pub fn link(&self, state: usize) -> usize {
47        self.states[state].link
48    }
49
50    pub fn number_of_substrings(&self) -> usize {
51        let mut total = 0;
52        for state in 1..self.state_size() {
53            let link = self.link(state);
54            total += self.length(state) - self.length(link);
55        }
56        total
57    }
58
59    fn push(&mut self, c: usize) {
60        let new_node = self.states.len();
61        let last = self.last;
62        self.states.push(State {
63            next: HashMap::new(),
64            link: !0,
65            len: self.states[last].len + 1,
66        });
67        let mut p = last;
68        while p != !0 && !self.states[p].next.contains_key(&c) {
69            self.states[p].next.insert(c, new_node);
70            p = self.states[p].link;
71        }
72        let q = if p == !0 { 0 } else { self.states[p].next[&c] };
73        if p == !0 || self.states[p].len + 1 == self.states[q].len {
74            self.states[new_node].link = q;
75        } else {
76            let new_q = self.states.len();
77            self.states.push(State {
78                next: self.states[q].next.clone(),
79                link: self.states[q].link,
80                len: self.states[p].len + 1,
81            });
82            self.states[q].link = new_q;
83            self.states[new_node].link = new_q;
84            while p != !0 && self.states[p].next[&c] == q {
85                self.states[p].next.insert(c, new_q);
86                p = self.states[p].link;
87            }
88        }
89        self.last = new_node;
90    }
91}
92
93impl FromIterator<usize> for SuffixAutomaton {
94    fn from_iter<I: IntoIterator<Item = usize>>(iter: I) -> Self {
95        let mut sa = SuffixAutomaton::new();
96        sa.extend(iter);
97        sa
98    }
99}
100
101impl Extend<usize> for SuffixAutomaton {
102    fn extend<T: IntoIterator<Item = usize>>(&mut self, iter: T) {
103        for c in iter {
104            self.push(c);
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::tools::Xorshift;
113
114    #[test]
115    fn test_suffix_automaton() {
116        let mut rng = Xorshift::default();
117        for _ in 0..100 {
118            let csize = rng.random(1usize..=10);
119            let n = rng.random(1usize..=100);
120            let s: Vec<usize> = rng.random_iter(0usize..csize).take(n).collect();
121            let sa = SuffixAutomaton::from_iter(s.iter().cloned());
122            let mut len = vec![0; sa.state_size()];
123
124            for i in 0..n {
125                let mut state = 0;
126                for (j, &c) in s[i..].iter().enumerate() {
127                    assert!(sa.transitions(state).contains_key(&c));
128                    state = sa.transitions(state)[&c];
129                    len[state] = len[state].max(j + 1);
130                }
131            }
132            for (state, &len) in len.iter().enumerate() {
133                assert_eq!(sa.length(state), len);
134            }
135            for state in 1..sa.state_size() {
136                assert_ne!(sa.link(state), !0);
137            }
138        }
139    }
140
141    #[test]
142    fn test_number_of_substrings() {
143        let mut rng = Xorshift::default();
144        for _ in 0..100 {
145            let csize = rng.random(1usize..=10);
146            let n = rng.random(1usize..=100);
147            let s: Vec<usize> = rng.random_iter(0usize..csize).take(n).collect();
148            let sa = SuffixAutomaton::from_iter(s.iter().cloned());
149            let mut substrings = std::collections::HashSet::new();
150            for i in 0..n {
151                for j in i + 1..=n {
152                    substrings.insert(&s[i..j]);
153                }
154            }
155            assert_eq!(sa.number_of_substrings(), substrings.len());
156        }
157    }
158}
competitive/string/suffix_automaton.rs

competitive/string/
suffix_automaton.rs
