github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/trie/ac/ac.go (about)

     1  package ac
     2  
     3  type ac struct {
     4  	root *acNode
     5  }
     6  
     7  type acNode struct {
     8  	fail *acNode
     9  	mark string
    10  	node map[rune]*acNode
    11  }
    12  
    13  func (a *ac) search(str string) []string {
    14  	p := a.root
    15  	resp := []string{}
    16  
    17  	for _, c := range str {
    18  		_, ok := p.node[c]
    19  		for !ok && p != a.root {
    20  			p = p.fail
    21  			_, ok = p.node[c]
    22  		}
    23  
    24  		if _, ok = p.node[c]; ok {
    25  			p = p.node[c]
    26  			if p.mark != "" {
    27  				resp = append(resp, p.mark)
    28  			}
    29  			if p.fail.mark != "" {
    30  				resp = append(resp, p.fail.mark)
    31  			}
    32  		}
    33  	}
    34  	return resp
    35  }
    36  
    37  func (a *ac) searchLongest(str string) string {
    38  	p := a.root
    39  	resp := ""
    40  
    41  	for _, c := range str {
    42  		_, ok := p.node[c]
    43  		for !ok && p != a.root {
    44  			p = p.fail
    45  			_, ok = p.node[c]
    46  		}
    47  
    48  		if _, ok = p.node[c]; ok {
    49  
    50  			p = p.node[c]
    51  			if p.mark != "" {
    52  				resp = p.mark
    53  			}
    54  			// if p.fail.mark != "" {
    55  			// resp = p.fail.mark
    56  			// }
    57  		}
    58  	}
    59  
    60  	return resp
    61  }
    62  
    63  func (a *ac) Insert(str string) {
    64  	p := a.root
    65  	for _, c := range str {
    66  		if p.node[c] == nil {
    67  			p.node[c] = &acNode{node: make(map[rune]*acNode)}
    68  		}
    69  		p = p.node[c]
    70  	}
    71  	p.mark = str
    72  }
    73  
    74  func (a *ac) BuildFail() {
    75  	que := newQueue()
    76  	que.push(&queueElem{p: a.root, n: a.root})
    77  
    78  	for que.size() != 0 {
    79  		z := que.pop()
    80  		z.n.fail = a.findFail(z.p, z.b)
    81  
    82  		for k, v := range z.n.node {
    83  			que.push(&queueElem{
    84  				p: z.n,
    85  				n: v,
    86  				b: k,
    87  			})
    88  		}
    89  	}
    90  }
    91  
    92  func (a *ac) findFail(parent *acNode, b rune) *acNode {
    93  	if parent == a.root {
    94  		return parent
    95  	}
    96  	if i, ok := parent.fail.node[b]; ok {
    97  		return i
    98  	}
    99  	return a.findFail(parent.fail, b)
   100  }
   101  
   102  type queueElem struct {
   103  	p *acNode
   104  	n *acNode
   105  	b rune
   106  }
   107  type queue struct {
   108  	qu []*queueElem
   109  }
   110  
   111  func newQueue() *queue {
   112  	return &queue{
   113  		qu: []*queueElem{},
   114  	}
   115  }
   116  
   117  func (q *queue) pop() *queueElem {
   118  	if len(q.qu) == 0 {
   119  		return &queueElem{}
   120  	}
   121  	x := q.qu[0]
   122  	q.qu = q.qu[1:]
   123  	return x
   124  }
   125  
   126  func (q *queue) push(x *queueElem) {
   127  	q.qu = append(q.qu, x)
   128  }
   129  
   130  func (q *queue) size() int {
   131  	return len(q.qu)
   132  }
   133  func NewAC() *ac {
   134  	r := &acNode{node: make(map[rune]*acNode)}
   135  	r.fail = r
   136  
   137  	return &ac{root: r}
   138  }