suffix_automaton/
lib.rs

1//! Suffix Automaton (SAM).
2//!
3//! A suffix automaton is a directed acyclic automaton that recognises
4//! exactly the set of substrings of a given string. It has at most
5//! `2n - 1` states for an input of length `n`, and every state
6//! corresponds to a distinct set of end positions in the original
7//! string — hence to a distinct set of substrings.
8//!
9//! Once built, you can answer:
10//!
11//! * Is `p` a substring of `s`? (O(|p|))
12//! * How many distinct substrings does `s` have? (sum over states of
13//!   `len(state) - len(link(state))`)
14//! * Number of occurrences of a pattern (when `state.occ` is computed)
15//! * Longest common substring between two strings
16//!
17//! # Example
18//!
19//! ```
20//! use suffix_automaton::SuffixAutomaton;
21//!
22//! let mut sam = SuffixAutomaton::new();
23//! for ch in "ababa".chars() {
24//!     sam.extend(ch);
25//! }
26//! assert!(sam.contains("aba"));
27//! assert!(sam.contains("bab"));
28//! assert!(!sam.contains("bb"));
29//! ```
30
31/// State of a suffix automaton.
32#[derive(Debug, Clone)]
33pub struct State {
34    /// Length of the longest substring represented by this state.
35    pub len: usize,
36    /// Suffix link: points to the state representing the longest proper
37    /// suffix of this state's substrings.
38    pub link: Option<usize>,
39    /// Transitions keyed by character.
40    pub next: std::collections::BTreeMap<char, usize>,
41}
42
43/// Suffix Automaton over a sequence of `char`s.
44#[derive(Debug, Clone)]
45pub struct SuffixAutomaton {
46    /// All states, indexed by id. State 0 is the initial state.
47    pub states: Vec<State>,
48    /// Id of the state representing the entire current string.
49    pub last: usize,
50}
51
52impl SuffixAutomaton {
53    /// Build an empty suffix automaton.
54    pub fn new() -> Self {
55        SuffixAutomaton {
56            states: vec![State {
57                len: 0,
58                link: None,
59                next: std::collections::BTreeMap::new(),
60            }],
61            last: 0,
62        }
63    }
64
65    /// Extend the automaton with a single character.
66    pub fn extend(&mut self, ch: char) {
67        let cur = self.states.len();
68        self.states.push(State {
69            len: self.states[self.last].len + 1,
70            link: None,
71            next: std::collections::BTreeMap::new(),
72        });
73
74        let mut p = self.last;
75        while p != usize::MAX && !self.states[p].next.contains_key(&ch) {
76            self.states[p].next.insert(ch, cur);
77            p = self.states[p].link.unwrap_or(usize::MAX);
78        }
79
80        if p == usize::MAX {
81            self.states[cur].link = Some(0);
82        } else {
83            let q = self.states[p].next[&ch];
84            if self.states[p].len + 1 == self.states[q].len {
85                self.states[cur].link = Some(q);
86            } else {
87                // Clone state q into a new state `clone`.
88                let clone = self.states.len();
89                self.states.push(State {
90                    len: self.states[p].len + 1,
91                    link: self.states[q].link,
92                    next: self.states[q].next.clone(),
93                });
94                while p != usize::MAX && self.states[p].next.get(&ch) == Some(&q) {
95                    self.states[p].next.insert(ch, clone);
96                    p = self.states[p].link.unwrap_or(usize::MAX);
97                }
98                self.states[q].link = Some(clone);
99                self.states[cur].link = Some(clone);
100            }
101        }
102
103        self.last = cur;
104    }
105
106    /// Build the automaton from an entire string in one call.
107    pub fn from_str(s: &str) -> Self {
108        let mut sam = SuffixAutomaton::new();
109        for ch in s.chars() {
110            sam.extend(ch);
111        }
112        sam
113    }
114
115    /// Returns `true` iff `pattern` is a substring of the built string.
116    pub fn contains(&self, pattern: &str) -> bool {
117        let mut v = 0usize;
118        for ch in pattern.chars() {
119            match self.states[v].next.get(&ch) {
120                Some(&u) => v = u,
121                None => return false,
122            }
123        }
124        true
125    }
126
127    /// Number of distinct substrings of the built string.
128    pub fn distinct_substring_count(&self) -> u64 {
129        // Σ (len(v) - len(link(v))) over all states v != 0.
130        let mut total: u64 = 0;
131        for (i, s) in self.states.iter().enumerate() {
132            if i == 0 {
133                continue;
134            }
135            let link_len = s.link.map(|l| self.states[l].len).unwrap_or(0);
136            total += (s.len - link_len) as u64;
137        }
138        total
139    }
140}
141
142impl Default for SuffixAutomaton {
143    fn default() -> Self {
144        Self::new()
145    }
146}
suffix_automaton/lib.rs

suffix_automaton/
lib.rs
