kmp_search can now find *all* indices, not just the first one

jabagawee · jabagawee · commit a5dc90cb919d · 2012-10-02T17:50:25.000-07:00
diff --git a/algorithms/searching/kmp_search.py b/algorithms/searching/kmp_search.py
@@ -15,28 +15,30 @@
 
     Psuedo Code: CLRS. Introduction to Algorithms. 3rd ed.
 
-    kmp_search.search(sorted_list) -> integer
-    kmp_search.search(sorted_list) -> False
+    kmp_search.search(sorted_list) -> list[integers]
+    kmp_search.search(sorted_list) -> list[empty]
 """
 
 
 def search(string, word):
     word_length = len(word)
     prefix = compute_prefix(word)
-    q = 0
+    offsets = []
+    q = 0 # q is the number of characters matched
     for i in xrange(len(string)):
         while q > 0 and word[q] != string[i]:
-            q = prefix[q - 1]
+            q = prefix[q - 1] # next character does not match
         if word[q] == string[i]:
-            q = q + 1
+            q += 1
         if q == word_length:
-            return i - word_length + 1
-    return False
+            offsets.append(i - word_length + 1)
+            q = prefix[q - 1] # look for next match
+    return offsets
 
 
 def compute_prefix(word):
     word_length = len(word)
-    prefix = [0] * word_length
+    prefix = [0 for i in xrange(word_length)]
     k = 0
 
     for q in xrange(1, word_length):
diff --git a/algorithms/tests/test_searching.py b/algorithms/tests/test_searching.py
@@ -41,7 +41,7 @@ def test_kmpsearch(self):
         self.string = "ABCDE FG ABCDEABCDEF"
         rv1 = kmp_search.search(self.string, "ABCDEA")
         rv2 = kmp_search.search(self.string, "ABCDER")
-        self.assertIs(rv1, 9)
+        self.assertIs(rv1[0], 9)
         self.assertFalse(rv2)
 
 
