From 59cf685395258ac6ad59e123e1abdc15d33c9bc6 Mon Sep 17 00:00:00 2001
From: June Park <junepark202012@gmail.com>
Date: Tue, 3 Sep 2024 00:49:53 +0900
Subject: [PATCH] fix(trie): trie returns invalid string when only the fist
 character matches, and the second doesn't match

---
 src/data-structures/trie/trie.ts            |  7 ++++++-
 test/unit/data-structures/trie/trie.test.ts | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/data-structures/trie/trie.ts b/src/data-structures/trie/trie.ts
index 65503f8..05c23dc 100644
--- a/src/data-structures/trie/trie.ts
+++ b/src/data-structures/trie/trie.ts
@@ -454,7 +454,12 @@ export class Trie<R = any> extends IterableElementBase<string, R, Trie<R>> {
     if (prefix) {
       for (const c of prefix) {
         const nodeC = startNode.children.get(c);
-        if (nodeC) startNode = nodeC;
+        if (nodeC) {
+          startNode = nodeC;
+        } else {
+          // Early return if the whole prefix is not found
+          return [];
+        }
       }
     }
 
diff --git a/test/unit/data-structures/trie/trie.test.ts b/test/unit/data-structures/trie/trie.test.ts
index 59e02f9..4b87714 100644
--- a/test/unit/data-structures/trie/trie.test.ts
+++ b/test/unit/data-structures/trie/trie.test.ts
@@ -836,6 +836,22 @@ describe('Trie operations', () => {
     expect(words).toEqual(['apple', 'appetizer']);
   });
 
+  it('Get no words when prefix not found, with no match from the first character', () => {
+    trie.add('apple');
+    trie.add('appetizer');
+    trie.add('banana');
+    const words = trie.getWords('cd');
+    expect(words).toEqual([]);
+  });
+
+  it('Get no words when prefix not found, with no match from the second character', () => {
+    trie.add('apple');
+    trie.add('appetizer');
+    trie.add('banana');
+    const words = trie.getWords('ab');
+    expect(words).toEqual([]);
+  });
+
   it('Tree Height', () => {
     trie.add('apple');
     trie.add('banana');