Search a string from a list of strings using a Trie.

task.py

"""

Lets say you have an array with a list of sub strings. Write a function that takes an input string and return true if the string can be divided into substrings contained in the array.
Example: array: [‘and’, ‘the’, ‘back’, ‘front’, ‘ory’ , ‘end’, ‘po’ , ‘pu’, ‘lar’]
function subDivideString(inputString):

if inputString ==’backend’   return true since array has ‘back’ , ‘end’
if inputString == ‘popular’  return true since array has ‘po’ , ‘pu’, ‘lar’
if inputString == ‘backwards’   return false since array doesnt have ‘wards’
if inputString == ‘swapan’   return false since array doesnt have any part of the word.
"""

from nose.tools import assert_equal

class Node(object):
    def __init__(self, is_end=False):
        self.children = {}
        self.is_end = is_end
        
class Trie(object):
    def __init__(self):
        self.root = Node()
        
    
    def insert(self, word):
        """
        Inser a word into the trie.
        """
        if not isinstance(word, str):
            raise TypeError
            
        current = self.root
        for ch in word:
            if ch not in current.children:
                current.children[ch] = Node()
            current = current.children[ch]
        current.is_end = True
          
        
def search_trie(list_of_words, data):
    """Search in the trie for the word."""
    if data is None or len(data) == 0:
        return False
    else:
        tr = Trie()
        for word in list_of_words:
            tr.insert(word)
            
        current = tr.root
        for x in data:
            if current.is_end:
                current = tr.root
            if x not in current.children:
                return False
            else:
                current = current.children[x]
        
        return current.is_end
            

        
def search_substrings(list_of_words, data):
    tr = Trie()
    for word in list_of_words:
        tr.insert(word)
        
    print search_trie(tr, data)
        
        
class TestTrie(object):
    def test_search(self, solution):
        arr = ["and", "the", "back", "front", "ory" , "end", "po" , "pu", "lar"]
        assert_equal(solution(arr, 'backend'), True)
        assert_equal(solution(arr, 'back'), True)
        assert_equal(solution(arr, 'theend'), True)
        assert_equal(solution(arr, 'popular'), True)
        assert_equal(solution(arr, 'pop'), False)
        assert_equal(solution(arr, ''), False)
        assert_equal(solution(arr, 'backwards'), False)
        assert_equal(solution(arr, 'swapan'), False)
        assert_equal(solution(arr, 'andthebackfrontoryendpopular'), True)
        print 'Success: test_trie_search'


def main():
    test = TestTrie()
    test.test_search(search_trie)
    

if __name__ == '__main__':
    main()