Word Prefixes and Tries: Solution

class Trie(object):  # ...
    def find_prefix_words(self, prefix):
        """Find words with a prefix.

        :prefix: string of prefix

        - Navigates through the trie for each letter in the prefix
        - Returns a list of each word that uses this prefix, making sure
          to append the prefix to the returned list of words.

              a                   b
        /-----|-----\            / \
       *      c      t          a   e
             / \    / \        /   / \
            e   t  *   e      t   *   t
            |   |      |      |       |
            *   *      *      *       *

            >>> trie = Trie(["a", "ace", "act", "ate", "bat", "be", "bet"])

        If we provide an empty string for the prefix, this returns
        all child words of the entire trie:

            >>> trie.find_prefix_words('')
            ['a', 'ace', 'act', 'ate', 'bat', 'be', 'bet']

        Otherwise, it navigates to the end of that prefix, and finds
        all child words:

            >>> trie.find_prefix_words('a')
            ['a', 'ace', 'act', 'ate']

            >>> trie.find_prefix_words('ac')
            ['ace', 'act']

            >>> trie.find_prefix_words('ace')
            ['ace']

        (note that they must be valid words; 'ac' does not appear, as it is
        not a word -- there's no complete-marker in the children of 'a'->'c')

        If the prefix can't be found, it returns no found words:

            >>> trie.find_prefix_words('z')
            []

        Let's add a node with *no* children:

            >>> trie.root.children.append(TrieNode('z'))

        That's not very helpful -- this isn't a complete word, and it has no
        other children. Some people might even say our trie is no longer valid.
        Let's make sure this doesn't break things, though -- it should
        truthfully return no-words-with-that-prefix still:

            >>> trie.find_prefix_words('z')
            []
        """

        # Find prefix

        # START SOLUTION

        node = self.root

        for letter in prefix:
            found = False
            for child in node.children:
                if child is not TrieNode.COMPLETE_WORD and child.data == letter:
                    found = child
                    break
            if not found:
                # Can't find prefix, so no words use it
                return []

            node = found

        # END SOLUTION

        # Get list of suffixes from this point downward. This doesn't
        # include the prefix itself -- so four our sample trie, above,
        # 'ac' -> [e', 't']

        # START SOLUTION

        suffixes = node.find_child_words()

        # END SOLUTION

        # Return list of words, joining the prefix to each suffix

        # START SOLUTION

        return [prefix + suffix for suffix in suffixes]