Clean python code with Trie


  • 7
    class TrieNode(object):
        def __init__(self):
            self.word = None
            self.children = {}
    
    
    class Trie(object):
        def __init__(self):
            self.root = TrieNode()
    
        def insert(self, word):
            root = self.root
            for char in word:
                root = root.children.setdefault(char, TrieNode())
            root.word = word
    
    
    class Solution(object):
        def search(self, i, j, root, board, m, n, r):
            char = board[i][j]
            if not (char and char in root.children):
                return
    
            board[i][j], root = None, root.children[char]
    
            if root.word:
                r.append(root.word)
                root.word = None
    
            for x, y in ((0, -1), (-1, 0), (0, 1), (1, 0)):
                ii, jj = i + x, j + y
                if 0 <= ii < m and 0 <= jj < n:
                    self.search(ii, jj, root, board, m, n, r)
    
            board[i][j] = char
    
        def findWords(self, board, words):
            if not board:
                return []
    
            tree = Trie()
            [tree.insert(word) for word in words]
    
            m, n, r = len(board), len(board[0]), []
    
            for i, row in enumerate(board):
                for j, char in enumerate(row):
                    self.search(i, j, tree.root, board, m, n, r)
            return r
    

    Thanks @julien6 for pointing out a corner-case error.


  • 1
    J

    Thanks for sharing! Neat code as usual :)

    Although this time I got skeptical about the:

    del oldRoot.children[char]
    

    So I've tested your code with a corner case (with leetcode judging), this does not give the expected answer:

    ["baabab","abaaaa","abaaab","ababba","aabbab","aabbba","aabaab"]
    ["aab", "bbaabaabaaaaabaababaaaaababb","aabbaaabaaabaabaaaaaabbaaaba","babaababbbbbbbaabaababaabaaa","bbbaaabaabbaaababababbbbbaaa","babbabbbbaabbabaaaaaabbbaaab","bbbababbbbbbbababbabbbbbabaa","babababbababaabbbbabbbbabbba","abbbbbbaabaaabaaababaabbabba","aabaabababbbbbbababbbababbaa","aabbbbabbaababaaaabababbaaba","ababaababaaabbabbaabbaabbaba","abaabbbaaaaababbbaaaaabbbaab","aabbabaabaabbabababaaabbbaab","baaabaaaabbabaaabaabababaaaa","aaabbabaaaababbabbaabbaabbaa","aaabaaaaabaabbabaabbbbaabaaa","abbaabbaaaabbaababababbaabbb","baabaababbbbaaaabaaabbababbb","aabaababbaababbaaabaabababab","abbaaabbaabaabaabbbbaabbbbbb","aaababaabbaaabbbaaabbabbabab","bbababbbabbbbabbbbabbbbbabaa","abbbaabbbaaababbbababbababba","bbbbbbbabbbababbabaabababaab","aaaababaabbbbabaaaaabaaaaabb","bbaaabbbbabbaaabbaabbabbaaba","aabaabbbbaabaabbabaabababaaa","abbababbbaababaabbababababbb","aabbbabbaaaababbbbabbababbbb","babbbaabababbbbbbbbbaabbabaa"]
    

    Out:

    > ["aab","abaabbbaaaaababbbaaaaabbbaab","ababaababaaabbabbaabbaabbaba"]
    

    Expected:

    > ['aab', 'abaabbbaaaaababbbaaaaabbbaab', 'ababaababaaabbabbaabbaabbaba', 'aabbbbabbaababaaaabababbaaba']
    

    Because the first two elements share the same tree path, and the first one (shorter) gets found first.
    Better just set the word to None.


  • 0

    Yeah, exactly! Thanks very much! I have updated the code.


  • 0

    Thanks @julien6, good catch! I have just added and published your test case.


  • 0
    S

    I did mine without a trie and the l33tcode server timed out on me. I then looked for a Trie implementation and found this. It's so clean and pretty :') well done :o


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.