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    use crate::tools::testutil::{exhaustive_sequences, sample_usize, structured_sequences};
114
115    #[test]
116    fn test_suffix_automaton() {
117        let mut rng = Xorshift::default();
118        let lengths = sample_usize(&mut rng, 16, 0..=100, 100);
119        for s in
120            exhaustive_sequences(0..3, 0..=8).chain(structured_sequences(&mut rng, 0..10, lengths))
121        {
122            let n = s.len();
123            let sa = SuffixAutomaton::from_iter(s.iter().cloned());
124            let mut len = vec![0; sa.state_size()];
125
126            for i in 0..n {
127                let mut state = 0;
128                for (j, &c) in s[i..].iter().enumerate() {
129                    assert!(sa.transitions(state).contains_key(&c));
130                    state = sa.transitions(state)[&c];
131                    len[state] = len[state].max(j + 1);
132                }
133            }
134            for (state, &len) in len.iter().enumerate() {
135                assert_eq!(sa.length(state), len);
136            }
137            for state in 1..sa.state_size() {
138                assert_ne!(sa.link(state), !0);
139            }
140        }
141    }
142
143    #[test]
144    fn test_number_of_substrings() {
145        let mut rng = Xorshift::default();
146        let lengths = sample_usize(&mut rng, 16, 0..=100, 100);
147        for s in
148            exhaustive_sequences(0..3, 0..=8).chain(structured_sequences(&mut rng, 0..10, lengths))
149        {
150            let n = s.len();
151            let sa = SuffixAutomaton::from_iter(s.iter().cloned());
152            let mut substrings = std::collections::HashSet::new();
153            for i in 0..n {
154                for j in i + 1..=n {
155                    substrings.insert(&s[i..j]);
156                }
157            }
158            assert_eq!(sa.number_of_substrings(), substrings.len());
159        }
160    }
161}
competitive/string/suffix_automaton.rs

competitive/string/
suffix_automaton.rs
