3.3.3 String Searching (Aho-Corasick)

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

/*

Given a set of strings (needles) and subsequent queries of texts (haystacks) to be searched,
determine all positions in which needles occur within the given haystacks in linear time using the
Aho-Corasick algorithm. The needles are arranged in a trie whose nodes carry failure links to the
longest proper suffix that is also a trie node. Scanning the haystack once while following
transitions and failure links then reports every match through each node's output list.

This implementation stores transitions in hash tables, which keeps lookups expected constant time
without assuming a fixed alphabet size. Output lists store both needles ending at each trie node and
needles inherited through its failure link.

- `AhoCorasick(needles)` constructs the finite-state automaton for a set of needle strings that are
  to be searched for subsequently in haystack queries. Empty needles are ignored.
- `find_all_in(haystack, report_match)` calls the function `report_match(s, pos)` once on each
  occurrence of each needle that occurs in the `haystack`, where `pos` is the starting position in
  the `haystack` at which string `s` (a matched needle) occurs. The matches will be reported in
  increasing order of their ending positions within the `haystack`.

Time Complexity:
- O(L + z) expected per call to the constructor, where $L$ is the total needle length and $z$ is the
  total number of inherited output entries.
- O(n + z) expected per call to `find_all_in(haystack, report_match)`, where $n$ is the length of
  `haystack` and $z$ is the number of matches.

Space Complexity:
- O(L + z) for storage of the automaton and inherited output entries.
- O(1) auxiliary for `find_all_in()`.

*/

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

class AhoCorasick {
  std::vector<string> needles;
  std::vector<int> fail;
  std::vector<std::unordered_map<char, int>> adj;
  std::vector<std::vector<int>> out;

  int next_state(int curr, char c) const {
    int next = curr;
    while (next != 0 && adj[next].find(c) == adj[next].end()) {
      next = fail[next];
    }
    auto it = adj[next].find(c);
    return (it != adj[next].end()) ? it->second : 0;
  }

 public:
  explicit AhoCorasick(const std::vector<string> &needles) : needles(needles) {
    int total_len = 0;
    for (const auto &needle : needles) {
      total_len += static_cast<int>(needle.size());
    }
    // The trie has at most total_len + 1 states: the root plus one per character when no prefixes
    // are shared. Sizing to total_len alone overflows for low-sharing needle sets (e.g. a single
    // needle, which still needs the root plus one node).
    int max_states = total_len + 1;
    fail.resize(max_states, -1);
    adj.resize(max_states);
    out.resize(max_states);
    int states = 1;
    for (int i = 0; i < static_cast<int>(needles.size()); i++) {
      if (needles[i].empty()) {
        continue;
      }
      int curr = 0;
      for (char c : needles[i]) {
        if (auto it = adj[curr].find(c); it != adj[curr].end()) {
          curr = it->second;
        } else {
          curr = adj[curr][c] = states++;
        }
      }
      out[curr].push_back(i);
    }
    std::queue<int> q;
    for (auto &[c, v] : adj[0]) {
      if (v != 0) {
        fail[v] = 0;
        q.push(v);
      }
    }
    while (!q.empty()) {
      int u = q.front();
      q.pop();
      for (auto &[c, v] : adj[u]) {
        int f = fail[u];
        while (f != 0 && adj[f].find(c) == adj[f].end()) {
          f = fail[f];
        }
        auto fit = adj[f].find(c);
        f = (fit != adj[f].end()) ? fit->second : 0;
        fail[v] = f;
        out[v].insert(out[v].end(), out[f].begin(), out[f].end());
        q.push(v);
      }
    }
  }

  template<typename Fn>
  void find_all_in(const string &haystack, Fn report_match) const {
    int state = 0;
    for (int i = 0; i < static_cast<int>(haystack.size()); i++) {
      state = next_state(state, haystack[i]);
      for (int idx : out[state]) {
        report_match(needles[idx], i - static_cast<int>(needles[idx].size()) + 1);
      }
    }
  }
};

/*** Example Usage ***/

#include <algorithm>
#include <cassert>
#include <utility>
using namespace std;

int main() {
  vector<string> needles;
  needles.push_back("a");
  needles.push_back("ab");
  needles.push_back("bab");
  needles.push_back("bc");
  needles.push_back("bca");
  needles.push_back("c");
  needles.push_back("caa");
  needles.push_back("abccab");

  AhoCorasick ac(needles);
  vector<pair<string, int>> matches;
  ac.find_all_in("abccab", [&](const string &needle, int pos) {
    matches.push_back({needle, pos});
  });
  sort(matches.begin(), matches.end());
  assert(
      (matches ==
       vector<pair<string, int>>{
           {"a", 0}, {"a", 4}, {"ab", 0}, {"ab", 4}, {"abccab", 0}, {"bc", 1}, {"c", 2}, {"c", 3}
       })
  );

  vector<pair<string, int>> none;
  ac.find_all_in("zzzz", [&](const string &needle, int pos) { none.push_back({needle, pos}); });
  assert(none.empty());
  return 0;
}