From 364763bcc7453275ed3cd0d33311d962a91b597c Mon Sep 17 00:00:00 2001
From: Syed Ibna Zubayear <47857423+Zubayear@users.noreply.github.com>
Date: Sat, 30 Aug 2025 22:19:28 +0600
Subject: [PATCH] feat: refactor Trie to O(1) Size, prevent duplicate
 insertions, lazy remove, and add edge-case tests

---
 structure/trie/trie.go            |  32 ++++++---
 structure/trie/trie_bench_test.go |  14 ++++
 structure/trie/trie_test.go       | 105 ++++++++++++++++++++++++++++++
 3 files changed, 142 insertions(+), 9 deletions(-)

diff --git a/structure/trie/trie.go b/structure/trie/trie.go
index 3665e2c3c..73cdfd6d3 100644
--- a/structure/trie/trie.go
+++ b/structure/trie/trie.go
@@ -7,6 +7,7 @@ package trie
 type Node struct {
 	children map[rune]*Node // map children nodes
 	isLeaf   bool           // current node value
+	total    int            // total words in the Trie
 }
 
 // NewNode creates a new Trie node with initialized
@@ -20,6 +21,9 @@ func NewNode() *Node {
 
 // insert a single word at a Trie node.
 func (n *Node) insert(s string) {
+	if len(s) == 0 {
+		return
+	}
 	curr := n
 	for _, c := range s {
 		next, ok := curr.children[c]
@@ -29,7 +33,10 @@ func (n *Node) insert(s string) {
 		}
 		curr = next
 	}
-	curr.isLeaf = true
+	if !curr.isLeaf {
+		curr.isLeaf = true
+		n.total++
+	}
 }
 
 // Insert zero, one or more words at a Trie node.
@@ -62,14 +69,15 @@ func (n *Node) Capacity() int {
 
 // Size returns the number of words in the Trie
 func (n *Node) Size() int {
-	r := 0
-	for _, c := range n.children {
-		r += c.Size()
-	}
-	if n.isLeaf {
-		r++
-	}
-	return r
+	//r := 0
+	//for _, c := range n.children {
+	//	r += c.Size()
+	//}
+	//if n.isLeaf {
+	//	r++
+	//}
+	//return r
+	return n.total
 }
 
 // remove lazily a word from the Trie node, no node is actually removed.
@@ -85,7 +93,13 @@ func (n *Node) remove(s string) {
 			return
 		}
 	}
+	if !next.isLeaf {
+		return
+	}
 	next.isLeaf = false
+	if n.total > 0 {
+		n.total--
+	}
 }
 
 // Remove zero, one or more words lazily from the Trie, no node is actually removed.
diff --git a/structure/trie/trie_bench_test.go b/structure/trie/trie_bench_test.go
index 7243a1bd9..22df0bc3a 100644
--- a/structure/trie/trie_bench_test.go
+++ b/structure/trie/trie_bench_test.go
@@ -79,3 +79,17 @@ func BenchmarkTrie_Remove_and_Compact(b *testing.B) {
 		n.Compact()
 	}
 }
+
+func BenchmarkTrie_Size(b *testing.B) {
+	insert := make([]string, 3000)
+	for i := 0; i < len(insert); i++ {
+		insert[i] = fmt.Sprintf("%f", rand.Float64())
+	}
+	n := NewNode()
+	n.Insert(insert...)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_ = n.Size()
+	}
+}
diff --git a/structure/trie/trie_test.go b/structure/trie/trie_test.go
index 7080b5898..71d8a83a8 100644
--- a/structure/trie/trie_test.go
+++ b/structure/trie/trie_test.go
@@ -125,6 +125,111 @@ func TestTrieRemove(t *testing.T) {
 	n.verifySizeCapa(t, 0, 1) // no words, only the root node left
 }
 
+// Test inserting empty string
+func TestTrieInsertEmptyString(t *testing.T) {
+	root := NewNode()
+	root.Insert("")
+
+	check := map[string]bool{
+		"": false,
+	}
+	root.verify(t, check)
+	root.verifySizeCapa(t, 0, root.Capacity())
+}
+
+// Test inserting words with shared prefix
+func TestTrieInsertSharedPrefix(t *testing.T) {
+	root := NewNode()
+	root.Insert("abc", "abcd", "abcde")
+
+	check := map[string]bool{
+		"abc":   true,
+		"abcd":  true,
+		"abcde": true,
+		"ab":    false,
+	}
+	root.verify(t, check)
+
+	root.verifySizeCapa(t, 3, root.Capacity())
+}
+
+// Test removing non-existent words
+func TestTrieRemoveNonExistentEdge(t *testing.T) {
+	root := NewNode()
+	root.Insert("abc", "def")
+
+	root.Remove("xyz") // does not exist
+
+	check := map[string]bool{
+		"abc": true,
+		"def": true,
+		"xyz": false,
+	}
+	root.verify(t, check)
+	root.verifySizeCapa(t, 2, root.Capacity())
+}
+
+// Test inserting duplicates multiple times
+func TestTrieInsertDuplicatesMultiple(t *testing.T) {
+	root := NewNode()
+	root.Insert("abc", "abc", "abc")
+
+	check := map[string]bool{
+		"abc": true,
+	}
+	root.verify(t, check)
+
+	// Size should count only once
+	root.verifySizeCapa(t, 1, root.Capacity())
+}
+
+// Test removing all words sequentially
+func TestTrieRemoveAllSequential(t *testing.T) {
+	root := NewNode()
+	root.Insert("a", "b", "c")
+
+	root.Remove("a")
+	root.Remove("b")
+	root.Remove("c")
+
+	check := map[string]bool{
+		"a": false,
+		"b": false,
+		"c": false,
+	}
+	root.verify(t, check)
+
+	// Size should be zero
+	root.verifySizeCapa(t, 0, root.Capacity())
+}
+
+// Test inserting substrings and removing middle one
+func TestTrieInsertSubstringsRemoveMiddle(t *testing.T) {
+	root := NewNode()
+	root.Insert("a", "ab", "abc")
+
+	root.Remove("ab")
+
+	check := map[string]bool{
+		"a":   true,
+		"ab":  false,
+		"abc": true,
+	}
+	root.verify(t, check)
+
+	// Size should be 2
+	root.verifySizeCapa(t, 2, root.Capacity())
+}
+
+// Test Compact on empty Trie
+func TestTrieCompactEmptyTrie(t *testing.T) {
+	root := NewNode()
+	if !root.Compact() {
+		t.Fatal("Empty Trie root should be removable after compaction")
+	}
+	root.verifySizeCapa(t, 0, 1) // only root remains
+}
+
 // --------------- helper functions ---------------------------
 
 // verify if provided words are present