3.5.4 Suffix Automaton

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180

/*

Maintains a compact deterministic automaton representing all substrings of a string. A suffix
automaton is useful for substring membership queries, counting distinct substrings, finding longest
common substrings, and many other string dynamic programming tasks.

Each state represents an equivalence class of substrings with the same set of ending positions. The
value `st[v].len` is the maximum length represented by state `v`, while `st[v].link` points to the
state representing the longest proper suffix of that class.

More precisely, state `v` contributes one distinct substring for each length greater than
`st[st[v].link].len` and at most `st[v].len`. When a character is appended, `append()` creates a
state for the new full string and walks suffix links from the previous `last`, adding missing
transitions to it. If an existing transition leads to a state whose maximum length is too large, a
shorter clone of that state is inserted and the affected transitions and suffix links are redirected
to it, splitting the ending-position classes without changing the recognized substrings. Summing
these length ranges gives the number of distinct substrings. Membership queries follow transitions,
while longest common substring queries retreat through suffix links when a match fails.

- `SuffixAutomaton()` constructs an empty automaton.
- `SuffixAutomaton(s)` constructs the automaton for string `s`.
- `append(c)` appends character `c` to the current string.
- `contains(t)` returns whether string `t` occurs as a substring.
- `first_occurrence(t)` returns the starting index of the first occurrence of `t`, or $-1$ if
  absent.
- `count_distinct_substrings()` returns the number of distinct nonempty substrings.
- `longest_common_substring(t)` returns one longest substring common to the built string and string
  `t`.

Time Complexity:
- O(n) expected per call to the constructor for a string of length $n$.
- O(1) expected amortized per call to `append()`.
- O(m) expected per call to `contains(t)`, `first_occurrence(t)`, or `longest_common_substring(t)`,
  where $m$ is the length of `t`.
- O(n) per call to `count_distinct_substrings()`.

Space Complexity:
- O(n) transition storage and O(n) states.
- O(1) auxiliary per character processed.

*/

#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
using std::string;

class SuffixAutomaton {
  struct State {
    int len, link, first_pos;
    std::unordered_map<char, int> next;

    State() : len(0), link(-1), first_pos(-1) {}
  };

  std::vector<State> st;
  int last;

 public:
  SuffixAutomaton() : st(1), last(0) {}

  explicit SuffixAutomaton(const string &s) : SuffixAutomaton() {
    for (char c : s) {
      append(c);
    }
  }

  void append(char c) {
    int cur = static_cast<int>(st.size());
    st.emplace_back();
    st[cur].len = st[last].len + 1;
    st[cur].first_pos = st[cur].len - 1;
    int p = last;
    while (p != -1 && st[p].next.find(c) == st[p].next.end()) {
      st[p].next[c] = cur;
      p = st[p].link;
    }
    if (p == -1) {
      st[cur].link = 0;
    } else {
      int q = st[p].next[c];
      if (st[p].len + 1 == st[q].len) {
        st[cur].link = q;
      } else {
        int clone = static_cast<int>(st.size());
        st.push_back(st[q]);
        st[clone].len = st[p].len + 1;
        while (p != -1 && st[p].next[c] == q) {
          st[p].next[c] = clone;
          p = st[p].link;
        }
        st[q].link = st[cur].link = clone;
      }
    }
    last = cur;
  }

  bool contains(const string &t) const {
    int v = 0;
    for (char c : t) {
      if (auto it = st[v].next.find(c); it != st[v].next.end()) {
        v = it->second;
      } else {
        return false;
      }
    }
    return true;
  }

  // Returns the starting index of the first (leftmost) occurrence of t as a substring, or -1 if t
  // does not occur. first_pos stores the end index of the first occurrence of each state's class,
  // so subtracting the length of t recovers its start.
  int first_occurrence(const string &t) const {
    int v = 0;
    for (char c : t) {
      auto it = st[v].next.find(c);
      if (it == st[v].next.end()) {
        return -1;
      }
      v = it->second;
    }
    return st[v].first_pos - static_cast<int>(t.size()) + 1;
  }

  int64_t count_distinct_substrings() const {
    int64_t res = 0;
    for (int v = 1; v < static_cast<int>(st.size()); v++) {
      res += st[v].len - st[st[v].link].len;
    }
    return res;
  }

  string longest_common_substring(const string &t) const {
    int v = 0, len = 0, best = 0, best_pos = -1;
    for (int i = 0; i < static_cast<int>(t.size()); i++) {
      while (v && !st[v].next.count(t[i])) {
        v = st[v].link;
        len = st[v].len;
      }
      if (auto it = st[v].next.find(t[i]); it != st[v].next.end()) {
        v = it->second;
        len++;
      } else {
        v = 0;
        len = 0;
      }
      if (len > best) {
        best = len;
        best_pos = i;
      }
    }
    return best == 0 ? "" : t.substr(best_pos - best + 1, best);
  }
};

/*** Example Usage ***/

#include <cassert>

int main() {
  SuffixAutomaton sa("ababa");
  assert(sa.contains("aba"));
  assert(sa.contains("bab"));
  assert(!sa.contains("abba"));
  // "ababa" has 9 distinct nonempty substrings.
  assert(sa.count_distinct_substrings() == 9);
  // The longest substring shared by "ababa" and "zzbabab" is "baba".
  assert(sa.longest_common_substring("zzbabab") == "baba");
  assert(sa.first_occurrence("aba") == 0);
  assert(sa.first_occurrence("bab") == 1);
  assert(sa.first_occurrence("abba") == -1);

  SuffixAutomaton online;
  online.append('a');
  online.append('b');
  online.append('c');
  assert(online.contains("bc"));
  return 0;
}