aho-corasick.go

package main

import (
	"fmt"
)

type TrieNode struct {
	children map[rune]*TrieNode
	fail     *TrieNode
	output   []string
}

type AhoCorasick struct {
	root *TrieNode
}

func NewAhoCorasick() *AhoCorasick {
	return &AhoCorasick{
		root: &TrieNode{
			children: make(map[rune]*TrieNode),
			fail:     nil,
			output:   nil,
		},
	}
}

func (ac *AhoCorasick) AddPattern(pattern string) {
	current := ac.root

	for _, ch := range pattern {
		if current.children[ch] == nil {
			current.children[ch] = &TrieNode{
				children: make(map[rune]*TrieNode),
				fail:     nil,
				output:   nil,
			}
		}
		current = current.children[ch]
	}

	current.output = append(current.output, pattern)
}

func (ac *AhoCorasick) buildFailTransitions() {
	queue := make([]*TrieNode, 0)

	// Set the fail pointer of all level-1 nodes to the root node
	for _, child := range ac.root.children {
		child.fail = ac.root
		queue = append(queue, child)
	}

	for len(queue) > 0 {
		current := queue[0]
		queue = queue[1:]

		for ch, child := range current.children {
			failNode := current.fail

			for failNode != nil && failNode.children[ch] == nil {
				failNode = failNode.fail
			}

			if failNode != nil {
				child.fail = failNode.children[ch]
			} else {
				child.fail = ac.root
			}

			child.output = append(child.output, child.fail.output...)
			queue = append(queue, child)
		}
	}
}

func (ac *AhoCorasick) Search(text string) []string {
	result := make([]string, 0)
	current := ac.root

	for _, ch := range text {
		for current.children[ch] == nil && current != ac.root {
			current = current.fail
		}

		if current.children[ch] != nil {
			current = current.children[ch]
		}

		result = append(result, current.output...)
	}

	return result
}

func main() {
	ac := NewAhoCorasick()
	ac.AddPattern("he")
	ac.AddPattern("she")
	ac.AddPattern("his")
	ac.AddPattern("hers")

	ac.buildFailTransitions()

	text := "ushers"
	matches := ac.Search(text)

	fmt.Printf("Text: %s\nMatches: %v\n", text, matches)
}