github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/wordsearch/trie/trie.go (about)

     1  package trie
     2  
     3  import "unicode/utf8"
     4  
     5  type Iter struct {
     6  	At    *Node
     7  	Index int
     8  }
     9  
    10  func (it Iter) Valid() bool { return it.Index >= 0 }
    11  func (it Iter) Match() bool {
    12  	if it.Index == 0 {
    13  		return it.At.Term
    14  	}
    15  	return it.Index == len(it.At.Suffix)
    16  }
    17  
    18  func (it Iter) Walk(r rune) Iter {
    19  	if it.At.Suffix != "" {
    20  		sr, sz := utf8.DecodeRuneInString(it.At.Suffix[it.Index:])
    21  		if sr == r {
    22  			return Iter{it.At, it.Index + sz}
    23  		}
    24  		return Iter{nil, -1}
    25  	}
    26  
    27  	p, ok := it.At.Pos(r)
    28  	if !ok {
    29  		return Iter{nil, -1}
    30  	}
    31  	return Iter{&it.At.Edges[p], 0}
    32  }
    33  
    34  type Node struct {
    35  	Label  rune
    36  	Term   bool
    37  	Suffix string
    38  	Edges  []Node
    39  }
    40  
    41  func (n *Node) Insert(word string) {
    42  	if word == "" {
    43  		n.Term = true
    44  		return
    45  	}
    46  
    47  	if len(n.Edges) == 0 {
    48  		if n.Suffix == "" {
    49  			n.Suffix = word
    50  		} else {
    51  			n.insertEdge(n.Suffix)
    52  			n.insertEdge(word)
    53  			n.Suffix = ""
    54  		}
    55  	} else {
    56  		n.insertEdge(word)
    57  	}
    58  }
    59  
    60  func (n *Node) Contains(word string) bool {
    61  	it := Iter{n, 0}
    62  	if !it.Valid() {
    63  		return false
    64  	}
    65  	for _, r := range word {
    66  		it = it.Walk(r)
    67  		if !it.Valid() {
    68  			return false
    69  		}
    70  	}
    71  	return it.Match()
    72  }
    73  
    74  func (n *Node) insertEdge(word string) {
    75  	r, sz := utf8.DecodeRuneInString(word)
    76  	p, ok := n.Pos(r)
    77  	if ok {
    78  		n.Edges[p].Insert(word[sz:])
    79  	} else {
    80  		edge := Node{Label: r}
    81  		edge.Insert(word[sz:])
    82  
    83  		n.Edges = append(n.Edges, edge)
    84  		copy(n.Edges[p+1:], n.Edges[p:])
    85  		n.Edges[p] = edge
    86  	}
    87  }
    88  
    89  func (n *Node) Pos(label rune) (int, bool) {
    90  	// TODO: use linear scan
    91  	i, j := 0, len(n.Edges)
    92  	for i < j {
    93  		h := i + (j-i)/2
    94  		if n.Edges[h].Label < label {
    95  			i = h + 1
    96  		} else {
    97  			j = h
    98  		}
    99  	}
   100  	if i < 0 || len(n.Edges) <= i {
   101  		return i, false
   102  	}
   103  	return i, n.Edges[i].Label == label
   104  }