結果

問題 No.430 文字列検索
ユーザー 草苺奶昔草苺奶昔
提出日時 2023-02-19 02:55:41
言語 Go
(1.22.1)
結果
AC  
実行時間 20 ms / 2,000 ms
コード長 4,858 bytes
コンパイル時間 11,976 ms
コンパイル使用メモリ 230,748 KB
実行使用メモリ 9,856 KB
最終ジャッジ日時 2024-11-10 01:05:44
合計ジャッジ時間 11,709 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge3
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1 ms
6,816 KB
testcase_01 AC 20 ms
9,856 KB
testcase_02 AC 6 ms
5,248 KB
testcase_03 AC 7 ms
5,248 KB
testcase_04 AC 1 ms
5,248 KB
testcase_05 AC 1 ms
5,248 KB
testcase_06 AC 1 ms
5,248 KB
testcase_07 AC 1 ms
5,248 KB
testcase_08 AC 2 ms
5,248 KB
testcase_09 AC 2 ms
5,248 KB
testcase_10 AC 1 ms
5,248 KB
testcase_11 AC 15 ms
7,040 KB
testcase_12 AC 15 ms
7,936 KB
testcase_13 AC 15 ms
7,936 KB
testcase_14 AC 13 ms
6,656 KB
testcase_15 AC 12 ms
5,376 KB
testcase_16 AC 10 ms
5,376 KB
testcase_17 AC 9 ms
5,376 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

// https://nyaannyaan.github.io/library/string/aho-corasick.hpp

// API
// Add(s, id)
// Build(heavy)
// Match(s, heavy)
// MatchFrom(s, pos, heavy)
// Move(pos, char)
// Next(pos, j)
// Count(pos)
// Size()

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(true)

	match := aho.Match(s, true)
	res := 0
	for i := 0; i < len(match); i++ {
		res += match[i]
	}

	fmt.Fprintln(out, res)

}

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,
	}
}
0