Counting words with a given prefix

Author: volodimyr

2185.counting_words_with_a_given_prefix/prefix.py

@@ -0,0 +1,66 @@
+# 2185. Counting words with a given prefix
+# Topics: 'Array', 'String', 'String Matching'
+
+# You are given an array of strings words and a string pref.
+
+# Return the number of strings in words that contain pref as a prefix.
+
+# A prefix of a string s is any leading contiguous substring of s.
+
+ 
+
+# Example 1:
+
+# Input: words = ["pay","attention","practice","attend"], pref = "at"
+# Output: 2
+# Explanation: The 2 strings that contain "at" as a prefix are: "attention" and "attend".
+
+# Example 2:
+
+# Input: words = ["leetcode","win","loops","success"], pref = "code"
+# Output: 0
+# Explanation: There are no strings that contain "code" as a prefix.
+
+ 
+
+# Constraints:
+
+#     1 <= words.length <= 100
+#     1 <= words[i].length, pref.length <= 100
+#     words[i] and pref consist of lowercase English letters.
+
+from typing import List
+
+class Node:
+    def __init__(self):
+        self.children = {}
+        self.times = 1
+
+class Trie:
+    def __init__(self):
+        self.root = Node()
+    
+    def insert(self, word: str, length: int):
+        cur = self.root
+        for i in range (length):
+            if word[i] not in cur.children:
+                cur.children[word[i]] = Node()
+            else:
+                cur.children[word[i]].times += 1
+            cur = cur.children[word[i]]
+    
+    def prefix(self, pref: str) -> int:
+        cur = self.root
+        for char in pref:
+            if char not in cur.children:
+                return 0
+            cur = cur.children[char]
+        return cur.times
+
+class Solution:
+    def prefixCount(self, words: List[str], pref: str) -> int:
+        trie = Trie()
+        for word in words:
+            if len(word) >= len(pref):
+                trie.insert(word, len(pref))
+        return trie.prefix(pref)

2185.counting_words_with_a_given_prefix/test_prefix.py

@@ -0,0 +1,55 @@
+import unittest
+from prefix import Solution
+
+class TestPrefixCount(unittest.TestCase):
+    def setUp(self):
+        self.sol = Solution()
+
+    def test_example1(self):
+        words = ["pay", "attention", "practice", "attend"]
+        pref = "at"
+        self.assertEqual(self.sol.prefixCount(words, pref), 2)
+
+    def test_example2(self):
+        words = ["leetcode", "win", "loops", "success"]
+        pref = "code"
+        self.assertEqual(self.sol.prefixCount(words, pref), 0)
+
+    def test_single_match(self):
+        words = ["apple", "apricot", "banana"]
+        pref = "app"
+        self.assertEqual(self.sol.prefixCount(words, pref), 1)
+
+    def test_all_match(self):
+        words = ["a", "ab", "abc", "abcd"]
+        pref = "a"
+        self.assertEqual(self.sol.prefixCount(words, pref), 4)
+
+    def test_no_match(self):
+        words = ["cat", "dog", "fish"]
+        pref = "z"
+        self.assertEqual(self.sol.prefixCount(words, pref), 0)
+
+    def test_prefix_longer_than_words(self):
+        words = ["hi", "hello"]
+        pref = "hello!"
+        self.assertEqual(self.sol.prefixCount(words, pref), 0)
+
+    def test_duplicate_words(self):
+        words = ["at", "at", "attention"]
+        pref = "at"
+        self.assertEqual(self.sol.prefixCount(words, pref), 3)
+
+    def test_empty_list(self):
+        words = []
+        pref = "any"
+        self.assertEqual(self.sol.prefixCount(words, pref), 0)
+
+    def test_single_letter_prefix(self):
+        words = ["x", "xy", "xyz", "yy"]
+        pref = "x"
+        self.assertEqual(self.sol.prefixCount(words, pref), 3)
+
+
+if __name__ == "__main__":
+    unittest.main()