Re-implement suffix_tree function

Create a suffix tree from an input string. This implementation uses a stack to
control the flow of the program.

I expected this attempt to be easier than my first attempt, but surprisingly, it
was similarly difficult. It took me ~30-45 minutes to successfully implement
this function, and I'm still not pleased with the final result.

1 files changed, 29 insertions, 0 deletions

diff --git a/scratch/facebook/hard/suffix-tree.py b/scratch/facebook/hard/suffix-tree.py
index 600d06591bdb..782678fb822c 100644
--- a/scratch/facebook/hard/suffix-tree.py
+++ b/scratch/facebook/hard/suffix-tree.py
@@ -1,4 +1,5 @@
 import random
+from collections import deque
 
 def exists(pattern, tree):
     """
@@ -42,6 +43,34 @@ def suffix_tree(xs):
                 parent = children
     return root
 
+def suffix_tree(x):
+    """
+    Creates a suffix from the input string, `x`. This implementation uses a
+    stack.
+    """
+    result = [None, []]
+    q = deque()
+    for i in range(len(x)):
+        q.append((result, x[i:]))
+    while q:
+        parent, x = q.popleft()
+        s = []
+        s.append((parent, x))
+        while s:
+            parent, x = s.pop()
+            if not x:
+                continue
+            c, rest = x[0], x[1:]
+            grafted = False
+            for child in parent[1]:
+                if c == child[0]:
+                    s.append((child, rest))
+                    grafted = True
+            if not grafted:
+                child = [c, []]
+                parent[1].append(child)
+                s.append((child, rest))
+    return result[1]
 
 ################################################################################
 # Tests