Algorithm_templates/ahoCorasick.cpp at main · seo-bo/Algorithm_templates · GitHub

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
class ahoCorasick //https ://github.com/seo-bo/Algorithm_templates/blob/main/ahoCorasick.cpp
{
private:
	struct Node
	{
		map<char, Node*>child;
		Node* fail;
		vector<int> out;
		Node() { fail = nullptr; }
	};
	Node* root;
	bool built;
	void add(const string& word, int id)
	{
		Node* cur = root;
		for (auto& i : word)
		{
			if (cur->child.find(i) == cur->child.end())
			{
				cur->child[i] = new Node();
			}
			cur = cur->child[i];
		}
		cur->out.push_back(id);
	}
	void build_pattern()
	{
		queue<Node*>q;
		root->fail = root;
		for (auto& [a, b] : root->child)
		{
			b->fail = root;
			q.push(b);
		}
		while (!q.empty())
		{
			Node* cur = q.front();
			q.pop();
			for (auto& [a, b] : cur->child)
			{
				Node* f = cur->fail;
				while (f != root && f->child.find(a) == f->child.end())
				{
					f = f->fail;
				}
				if (f->child.find(a) == f->child.end())
				{
					b->fail = root;
				}
				else
				{
					b->fail = f->child[a];
				}
				b->out.insert(b->out.end(), b->fail->out.begin(), b->fail->out.end());
				q.push(b);
			}
		}
	}
	vector<pair<int, int>> search(const string& word)
	{
		vector<pair<int, int>>ans;
		Node* cur = root;
		for (int i = 0; i < (int)word.size(); ++i)
		{
			char ok = word[i];
			while (cur != root && cur->child.find(ok) == cur->child.end())
			{
				cur = cur->fail;
			}
			if (cur->child.find(ok) != cur->child.end())
			{
				cur = cur->child[ok];
			}
			for (auto& j : cur->out)
			{
				ans.push_back(make_pair(j, i));
			}
		}
		return ans;
	}
public:
	ahoCorasick() { root = new Node(); built = false; }
	void insert(const string& word, int id)
	{
		add(word, id);
	}
	vector<pair<int, int>> find(const string& word)
	{
		if (!built)
		{
			build_pattern();
			built = true;
		}
		return search(word);
	}
};