#841842 (Go) No.430 文字列検索

• 問題: No.430 文字列検索
• ユーザー: 草苺奶昔
• 提出日時: 2023-02-19 02:54:12
• 言語: Go (1.22.1)
• 結果: AC
• 実行時間: 16 ms / 2,000 ms
• コード長: 4,583 bytes
• コンパイル時間: 11,220 ms
• コンパイル使用メモリ: 213,400 KB
• 実行使用メモリ: 10,212 KB
• 最終ジャッジ日時: 2023-09-27 13:32:55
• 合計ジャッジ時間: 11,662 ms
• ジャッジサーバーID （参考情報）: judge11 / judge15

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 4,376 KB
testcase_01	AC	16 ms; 10,212 KB
testcase_02	AC	6 ms; 5,644 KB
testcase_03	AC	6 ms; 5,644 KB
testcase_04	AC	1 ms; 4,376 KB
testcase_05	AC	1 ms; 4,380 KB
testcase_06	AC	1 ms; 4,380 KB
testcase_07	AC	1 ms; 4,376 KB
testcase_08	AC	2 ms; 4,380 KB
testcase_09	AC	2 ms; 4,376 KB
testcase_10	AC	1 ms; 4,376 KB
testcase_11	AC	15 ms; 8,056 KB
testcase_12	AC	15 ms; 10,140 KB
testcase_13	AC	16 ms; 10,140 KB
testcase_14	AC	13 ms; 8,032 KB
testcase_15	AC	10 ms; 7,812 KB
testcase_16	AC	9 ms; 7,812 KB
testcase_17	AC	8 ms; 7,816 KB

ソースコード

package main

import (
	"bufio"
	"fmt"
	"os"
	"sort"
)

func main() {
	const INF int = int(1e18)
	const MOD int = 998244353

	in := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	var s string
	fmt.Fscan(in, &s)
	var n int
	fmt.Fscan(in, &n)
	aho := NewAhoCorasick(26, 'A')
	for i := 0; i < n; i++ {
		var t string
		fmt.Fscan(in, &t)
		aho.Add(t, i)
	}
	aho.Build(false)

	match := aho.Match(s, false)
	fmt.Println(match[0])
}

type AhoCorasick struct {
	*Trie
	count []int
}

// size: 字符集大小
// margin: 字符集起始字符
func NewAhoCorasick(size int, margin byte) *AhoCorasick {
	res := &AhoCorasick{Trie: NewTrie(size+1, margin)}
	return res
}

func (ac *AhoCorasick) Build(heavy bool) {
	n := len(ac.stack)
	ac.count = make([]int, n)
	for i := 0; i < n; i++ {
		if heavy {
			sort.Ints(ac.stack[i].indexes)
		}
		ac.count[i] = len(ac.stack[i].indexes)
	}

	var que []int
	for i := 0; i < ac.size-1; i++ {
		if *ac.Next(0, i) != -1 {
			*ac.Next(*ac.Next(0, i), ac.size-1) = 0
			que = append(que, *ac.Next(0, i))
		} else {
			*ac.Next(0, i) = 0
		}
	}

	for len(que) > 0 {
		x := ac.stack[que[0]]
		fail := x.next[ac.size-1]
		ac.count[que[0]] += ac.count[fail]
		que = que[1:]

		for i := 0; i < ac.size-1; i++ {
			nx := &x.next[i]
			if *nx < 0 {
				*nx = *ac.Next(fail, i)
				continue
			}
			que = append(que, *nx)
			*ac.Next(*nx, ac.size-1) = *ac.Next(fail, i)
			if heavy {
				idx := ac.stack[*nx].indexes
				idy := ac.stack[*ac.Next(fail, i)].indexes
				idz := make([]int, 0, len(idx)+len(idy))

				// set union
				i, j := 0, 0
				for i < len(idx) && j < len(idy) {
					if idx[i] < idy[j] {
						idz = append(idz, idx[i])
						i++
					} else if idx[i] > idy[j] {
						idz = append(idz, idy[j])
						j++
					} else {
						idz = append(idz, idx[i])
						i++
						j++
					}
				}
				for i < len(idx) {
					idz = append(idz, idx[i])
					i++
				}
				for j < len(idy) {
					idz = append(idz, idy[j])
					j++
				}

				ac.stack[*nx].indexes = idz
			}
		}

	}
}

func (ac *AhoCorasick) Match(s string, heavy bool) []int {
	size := 1
	if heavy {
		size = ac.Size()
	}
	res := make([]int, size)
	pos := 0
	for i := 0; i < len(s); i++ {
		pos = *ac.Next(pos, int(s[i]-ac.margin))
		if heavy {
			for _, x := range ac.stack[pos].indexes {
				res[x]++
			}
		} else {
			res[0] += ac.count[pos]
		}
	}
	return res
}

func (ac *AhoCorasick) MatchFrom(s string, pos int, heavy bool) []int {
	size := 1
	if heavy {
		size = ac.Size()
	}
	res := make([]int, size)
	for i := 0; i < len(s); i++ {
		pos = *ac.Next(pos, int(s[i]-ac.margin))
		if heavy {
			for _, x := range ac.stack[pos].indexes {
				res[x]++
			}
		} else {
			res[0] += ac.count[pos]
		}
	}
	return res
}

func (ac *AhoCorasick) Count(pos int) int {
	return ac.count[pos]
}

//
//
//
//
type Trie struct {
	size   int
	margin byte
	stack  []*trieNode
}

type trieNode struct {
	index   int // 最后一次被更新的字符串的索引
	key     byte
	indexes []int // 存储了哪些字符串的索引
	next    []int // children position
}

// size: 字符集大小
// margin: 字符集起始字符
func NewTrie(size int, margin byte) *Trie {
	res := &Trie{size: size, margin: margin}
	root := res.newNode('$')
	res.stack = append(res.stack, root)
	return res
}

func (t *Trie) Add(s string, index int) {
	pos := 0
	for i := 0; i < len(s); i++ {
		k := int(s[i] - t.margin)
		if *t.Next(pos, k) != -1 {
			pos = *t.Next(pos, k)
			continue
		}
		nextPos := len(t.stack)
		*t.Next(pos, k) = nextPos
		node := t.newNode(s[i])
		t.stack = append(t.stack, node)
		pos = nextPos
	}
	t.stack[pos].index = index
	t.stack[pos].indexes = append(t.stack[pos].indexes, index)
}

func (t *Trie) Find(s string) int {
	pos := 0
	for i := 0; i < len(s); i++ {
		k := int(s[i] - t.margin)
		if *t.Next(pos, k) == -1 {
			return -1
		}
		pos = *t.Next(pos, k)
	}
	return pos
}

func (t *Trie) Move(pos int, c byte) int {
	if pos < 0 || pos >= len(t.stack) {
		return -1
	}
	return *t.Next(pos, int(c-t.margin))
}

func (t *Trie) Index(pos int) int {
	if pos < 0 {
		return -1
	}
	return t.stack[pos].index
}

func (t *Trie) IndexAll(pos int) []int {
	if pos < 0 {
		return []int{}
	}
	return t.stack[pos].indexes
}

// Trie 中的节点数(包含根节点).
func (t *Trie) Size() int { return len(t.stack) }

func (t *Trie) Next(pos, j int) *int {
	return &t.stack[pos].next[j]
}

func (t *Trie) newNode(c byte) *trieNode {
	next := make([]int, t.size)
	for i := range next {
		next[i] = -1
	}
	return &trieNode{
		index: -1,
		key:   c,
		next:  next,
	}
}