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

     1  package trie
     2  
     3  type Uncompact struct {
     4  	root unode
     5  }
     6  
     7  type unode struct {
     8  	label  byte
     9  	term   bool
    10  	suffix string
    11  	edges  []unode
    12  }
    13  
    14  func (t *Uncompact) Insert(word string) {
    15  	t.root.insert(word)
    16  }
    17  
    18  func (n *unode) count() int {
    19  	t := 1
    20  	for i := range n.edges {
    21  		t += n.edges[i].count()
    22  	}
    23  	return t
    24  }
    25  
    26  func (n *unode) insert(word string) {
    27  	if word == "" {
    28  		n.term = true
    29  		return
    30  	}
    31  
    32  	if len(n.edges) == 0 {
    33  		if n.suffix == "" {
    34  			if safesuffix(word) {
    35  				n.suffix = word
    36  			} else {
    37  				n.insertEdge(word)
    38  			}
    39  		} else {
    40  			n.insertEdge(n.suffix)
    41  			n.insertEdge(word)
    42  			n.suffix = ""
    43  		}
    44  	} else {
    45  		n.insertEdge(word)
    46  	}
    47  }
    48  
    49  func (n *unode) insertEdge(word string) {
    50  	p, ok := n.pos(word[0])
    51  	if ok {
    52  		n.edges[p].insert(word[1:])
    53  	} else {
    54  		edge := unode{label: word[0]}
    55  		edge.insert(word[1:])
    56  
    57  		n.edges = append(n.edges, unode{})
    58  		copy(n.edges[p+1:], n.edges[p:])
    59  		n.edges[p] = edge
    60  	}
    61  }
    62  
    63  func (n *unode) pos(label byte) (int, bool) {
    64  	// TODO: use linear scan
    65  	i, j := 0, len(n.edges)
    66  	for i < j {
    67  		h := i + (j-i)/2
    68  		if n.edges[h].label < label {
    69  			i = h + 1
    70  		} else {
    71  			j = h
    72  		}
    73  	}
    74  	if i < 0 || len(n.edges) <= i {
    75  		return i, false
    76  	}
    77  	return i, n.edges[i].label == label
    78  }
    79  
    80  func suffixlen(a, b string) int {
    81  	n := len(a)
    82  	if len(b) < n {
    83  		n = len(b)
    84  	}
    85  	if maxsuffix < n {
    86  		n = maxsuffix
    87  	}
    88  
    89  	for i := 0; i < n; i++ {
    90  		if a[i] != b[i] {
    91  			return i
    92  		}
    93  	}
    94  
    95  	return n
    96  }