Aho-Corasick#

import "container/list"

type Result struct {
	index int
	pattern string
}

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

func newNode() *node {
	children := make(map[rune]*node)
	return &node{children: children, fail: nil, output: nil}
}

type AhoCorasick struct {
	root *node
}

func (a *AhoCorasick) Init(patterns []string) {
	a.root = newNode()
	for _, pattern := range patterns {
		a.insert(pattern)
	}
	a.connect()
}

func (a *AhoCorasick) insert(pattern string) {
	now := a.root
	for _, char := range []rune(pattern) {
		if _, exists := now.children[char]; !exists {
			now.children[char] = newNode()
		}
		now = now.children[char]
	}
	if now.output == nil {
		now.output = make([]string, 0)
	}
	now.output = append(now.output, pattern)
}

func (a *AhoCorasick) connect() {
	a.root.fail = a.root

	q := list.New()
	q.PushBack(a.root)

	for q.Len() > 0 {
		now := q.Remove(q.Front()).(*node)
		for key, next := range now.children {
			if now == a.root {
				next.fail = a.root
			} else {
				dst := now.fail
				for dst != a.root && dst.children[key] == nil {
					dst = dst.fail
				}
				if child, exists := dst.children[key]; exists {
					dst = child
				}
				next.fail = dst
			}
			next.output = append(next.output, next.fail.output...)
			q.PushBack(next)
		}
	}
}

func (a *AhoCorasick) search(word string) []Result {
	now := a.root
	res := make([]Result, 0)
	
	for i, char := range []rune(word) {
		for now != a.root && now.children[char] == nil {
			now = now.fail
		}
		
		if _, exists := now.children[char]; exists {
			now = now.children[char]
		}
		
		for _, p := range now.output {
			res = append(res, Result{i - len(p) + 1, p})
		}
	}
	return res
}

class Node:
	def __init__(self):
		self.children = dict()
		self.fail = None
		self.output = []


class AhoCorasick:
	def __init__(self, patterns: List[str]):
		self.root = Node()
		for pattern in patterns:
			self.insert(pattern)
		self.connect()

	def insert(self, pattern: str):
		'''
		아호-코라식 자료구조에 패턴을 삽입
		:param pattern: 패턴 문자열
		:type pattern: str
		'''
		now = self.root
		for char in pattern:
			now = now.children.setdefault(char, Node())
		now.output.append(pattern)

	def connect(self):
		'''
		실패 링크 연결
		'''
		from collections import deque
		
		self.root.fail = self.root

		q = deque([self.root])
		while q:
			now = q.popleft()
			for char in now.children:
				nxt = now.children[char]
				if now is self.root:
					nxt.fail = self.root
				else:
					dst = now.fail
					while dst is not self.root and char not in dst.children:
						dst = dst.fail
					if char in dst.children:
						dst = dst.children[char]
					nxt.fail = dst
				nxt.output += nxt.fail.output
				q.append(nxt)

	def search(self, word: str) -> List[Tuple[int, str]]:
		'''
		word에서 매칭되는 모든 패턴과, 해당 패턴이 시작하는 위치를 출력
		:param word: 검색 문자열
		:type word: str
		:return: (index, pattern)형태의 패턴과 해당 패턴의 시작 인덱스
		:rtype: List[Tuple[int, str]]
		'''
		now = self.root
		res = []
		for i, char in enumerate(word):
			while now is not self.root and char not in now.children:
				now = now.fail

			if char in now.children:
				now = now.children[char]
	
			for p in now.output:
				res.append((i - len(p) + 1, p))
		return res