[Feature] Added generating word squares.

Author: exterkamp

README.md

@@ -89,4 +89,5 @@ Different algorithmic programs.  Grouped by general topic.
   * [Reduce](algorithms/string/longest_common_prefix.py)
 * [Finding Substring](algorithms/string/rabin_karp_substring_search.py)
   * [Rabin-Karp](algorithms/string/rabin_karp_substring_search.py)
-* [Calculate Edit Distance](algorithms/string/edit_distance.py)
+* [Calculate Edit Distance](algorithms/string/edit_distance.py)
+* [Generate Word Squares](algorithms/string/word_squares.py)

algorithms/string/word_squares.py

@@ -0,0 +1,28 @@
+import structures
+from structures.trie import Trie
+
+def word_squares(words):
+
+    trie = Trie(words)
+
+    def construct(square):
+        if len(square) == len(square[0]):
+            construct.squares.append(square)
+            return
+        # find prefix, based on len of square
+        prefix = ""
+        idx = len(square)
+        for i in range(idx):
+            prefix += square[i][idx]
+        # get all prefix words
+        prefix_words = trie.get_all_common_prefix(prefix)
+        for prefix_word in prefix_words:
+            construct(square + [prefix_word])
+        
+    construct.squares = []
+
+    # start all squares
+    for word in words:
+        construct([word])
+
+    return construct.squares

structures/trie.py

@@ -14,4 +14,24 @@ def add(self, word):
         current = self.head
         for letter in word:
             current = current.setdefault(letter, {})
-        current[self.eof] = self.eof
+        current[self.eof] = self.eof
+    
+    def get_all_common_prefix(self, prefix):
+        idx = 0
+        node = self.head
+        while idx < len(prefix):
+            letter = prefix[idx]
+            if letter not in node:
+                return []
+            node = node[letter]
+            idx += 1
+        # get all words after this node
+        def getWords(node, prefix):
+            if self.eof in node:
+                return [prefix]
+            words = []
+            for key, val in node.items():
+                words.extend(getWords(val, prefix + key))
+            return words
+        
+        return getWords(node, prefix)

tests/test_word_squares.py

@@ -0,0 +1,13 @@
+import unittest
+from algorithms.string.word_squares import word_squares
+
+class TestWordSquares(unittest.TestCase):
+
+    def test_word_squares_basic(self):
+        self.assertEqual([['baba', 'abat', 'baba', 'atak']], word_squares(['baba','atak','abat','delt','quad']))
+
+    def test_words_squares_multiple_results(self):
+        self.assertEqual([['will', 'idle', 'lliw', 'lewd'], ['will', 'idle', 'llow', 'lewd']], word_squares(['will','idle','lliw','lewd','llow']))
+
+    def test_no_results(self):
+        self.assertEqual([], word_squares(['baba','free','quad','high','jill']))