Suffix Automaton

Suffix Automaton (SAM).

A suffix automaton is a directed acyclic automaton that recognises exactly the set of substrings of a given string. It has at most 2n - 1 states for an input of length n, and every state corresponds to a distinct set of end positions in the original string — hence to a distinct set of substrings.

Once built, you can answer:

  • Is p a substring of s? (O(|p|))

  • How many distinct substrings does s have? (sum over states of len(state) - len(link(state)))

  • Number of occurrences of a pattern (when state.occ is computed)

  • Longest common substring between two strings

The Rust crate lives at suffix_automaton/src/lib.rs; build & run it with:

$ cd suffix_automaton && cargo run
$ cd suffix_automaton && cargo test

Generated API reference for the Suffix Automaton crate. View the rendered API reference →

  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}