package main

import "fmt"

// 字典树节点
type Node struct {
    // 下一跳集合，由于题目中使用小写字母，因此使用固定大小的数组做字符映射，例如0->a, 1->b
    Nexts [26]*Node
    // 匹配的字符串集合，数组中的每一项表示从当前下标开始的前Exists[i]个为匹配的字符串
    Exists []int
    // 失效指针，用于匹配未命中时回退
    Fail *Node
}

// 向字典树中添加字符串
func (node *Node) Add(s string) {
    n := node
    for i := range s {
        // 得到当前字符
        c := s[i] - 'a'
        if n.Nexts[c] == nil {
            n.Nexts[c] = &Node{}
        }
        n = n.Nexts[c]
    }
    // 标记当前节点为匹配字符串的末端
    n.Exists = append(n.Exists, len(s))
}

// 构建符合AC自动机的字典树
func (node *Node) BuildAC() {
    root := node
    // 利用队列机制逐层遍历节点
    // 先将root的非空子节点加入队列中
    q := make([]*Node, 0)
    for _, next := range node.Nexts {
        if next != nil {
            // 初始的fail指针指向root
            next.Fail = root
            q = append(q, next)
        }
    }
    // root节点的fail指针指向本身
    root.Fail = root

    // BFS遍历
    for len(q) != 0 {
        top := q[0]
        q = q[1:]

        // 以当前节点视角决策下一跳的fail指针
        for c, next := range top.Nexts {
            // 跳过空对象
            if next == nil {
                continue
            }
            // 递归查找符合fail指针标准的节点，直至root节点
            fail := top.Fail
            for fail != root && fail.Nexts[c] == nil {
                fail = fail.Fail
            }

            // 如果找到
            if fail.Nexts[c] != nil {
                // 注意，这里需要指向失效指针的下一跳，因为由前面可知fail.Nexts[c] != nil，即fail.Nexts[c]就是当前next的等效备份
                next.Fail = fail.Nexts[c]
                // 如果找到的fail指针包含待匹配项，则追加
                if len(fail.Nexts[c].Exists) != 0 {
                    next.Exists = append(next.Exists, fail.Nexts[c].Exists...)
                }
            } else {
                // 默认指向root
                next.Fail = root
            }
            // 将下一跳加入队列
            q = append(q, next)
        }
    }
}

// 查找字符串是否是以匹配项为结尾
func (node *Node) Search(s []byte) bool {
    n := node
    root := node
    for i := range s {
        c := s[i] - 'a'
        // 递归查找匹配项，若下一跳未命中，则回退fail指针，直至root节点
        for n != root && n.Nexts[c] == nil {
            n = n.Fail
        }

        if n.Nexts[c] != nil {
            n = n.Nexts[c]
        }
    }
    // 至此为匹配到了字符串的最后一个字符，判断是否命中
    return len(n.Exists) != 0
}

// 查找字符串中包含的匹配项，返回string数组
func (node *Node) Query(s []byte) []string {
    n := node
    root := node
    ans := make([]string, 0)
    for i := range s {
        c := s[i] - 'a'
        // 递归查找匹配项，若下一跳未命中，则回退fail指针，直至root节点
        for n != root && n.Nexts[c] == nil {
            n = n.Fail
        }

        if n.Nexts[c] != nil {
            n = n.Nexts[c]
        }
        if len(n.Exists) != 0 {
            // 拼接匹配到的字符串
            for _, offset := range n.Exists {
                ans = append(ans, string(s[i+1-offset:i+1]))
            }
        }
    }
    return ans
}

type StreamChecker struct {
    root *Node
    s    []byte
}

func Constructor(words []string) StreamChecker {
    sc := StreamChecker{
        root: &Node{},
        s:    make([]byte, 0),
    }
    for _, w := range words {
        sc.root.Add(w)
    }
    sc.root.BuildAC()
    return sc
}

func (this *StreamChecker) Query(letter byte) bool {
    this.s = append(this.s, letter)
    end := len(this.s)
    start := end - 200
    if start < 0 {
        start = 0
    }
    return this.root.Search(this.s[start:end])
}

func main() {
    c := Constructor([]string{"his", "hers", "he", "she"})
    fmt.Println(c.root.Query([]byte("ahishers"))) // 输出[his she he hers]
    fmt.Println(c.Query('a')) // 输出false
    fmt.Println(c.Query('h')) // 输出false
    fmt.Println(c.Query('e')) // 输出true
    fmt.Println(c.Query('r')) // 输出false
    fmt.Println(c.Query('s')) // 输出true
}