jx/internal/search/trie.go

// Copyright 2024 Matthew Rich <matthewrich.conf@gmail.com>. All rights reserved.

package search

import (
)

var (
	AlphabetSize uint = 26
	FirstAlphabetChar = uint8('a')
)

type String string

type TrieNode struct {
	Children []*TrieNode
	Words int
}

func NewTrieNode() *TrieNode {
	return &TrieNode{ Children: make([]*TrieNode, AlphabetSize), Words: 0 }
}

func (r *TrieNode) Insert(key string) {
	current := r
	for _, c := range key {
		charIndex := uint8(c) - FirstAlphabetChar
		if current.Children[charIndex] == nil {
			current.Children[charIndex] = NewTrieNode()
		}
		current = current.Children[charIndex]
	}
	current.Words++
}

func (s String) Index(root *TrieNode) {
	root.Insert(string(s))
}

func (r *TrieNode) Search(key string) bool {
	current := r
	for _, c := range key {
		charIndex := uint8(c) - FirstAlphabetChar
		if current.Children[charIndex] == nil {
			return false
		}
		current = current.Children[charIndex]
	}
	return (current.Words > 0)
}

func (r *TrieNode) Delete(word string) bool {
	current := r
	var last *TrieNode = nil
	var lastChar rune
	for _, c := range word {
		charIndex := uint8(c) - FirstAlphabetChar
		if current.Children[charIndex] == nil {
			return false
		} else {
			count := 0
			for i := uint(0); i < AlphabetSize; i++ {
				if current.Children[i] != nil {
					count++
				}
			}
			if count > 1 {
				last = current
				lastChar = c
			}
			current = current.Children[charIndex]
		}
	}
	count := 0
	for i := uint(0); i < AlphabetSize; i++ {
		if current.Children[i] != nil {
			count++
		}
	}

	if (count > 0) { // is prefix
		current.Words--
		return true
	}

	if last != nil { // shares a common prefix with other words
		last.Children[lastChar] = nil
		return true
	} else { // does not share a common prefix
		r.Children[word[0]] = nil
		return true
	}
}