Find smallest subset prefixes

Question

Here is my code for this problem. Any bugs, performance in terms of algorithm time complexity, code style advice are appreciated.

The areas which I want to improve, but do not know how to improve are:

1. I sort the strings by length and I'm not sure if it's a smarter idea not using sorting at all.
2. I used a few set operations, which looks a bit ugly. I want to learn smarter ideas.

Problem:

Given a set of strings, return the smallest subset that contains the longest possible prefixes for every string.

If the list is ['foo', 'foog', 'food', 'asdf'] return ['foo', 'asdf']

The return is foo since foo is prefix for foo (itself), prefix for foog and prefix for food (in other words, foo could "represent" longer string like foog and food). Output also contains asdf because it is not prefix for any other words in the input list, so output itself.

The empty set is not a correct answer because it does not contain the longest possible prefixes.

Source code:

from collections import defaultdict
class TrieTreeNode:
    def __init__(self):
        self.children = defaultdict(TrieTreeNode)
        self.isEnd = False
        self.count = 0
    def insert(self, word, result_set):
        node = self
        for i,w in enumerate(word):
            node = node.children[w]
            if node.isEnd:
                result_set[word[:i+1]].add(word)
        node.isEnd = True

if __name__ == "__main__":
    words = ['foo', 'foog', 'food', 'asdf']
    words_length = []
    sorted_words = []
    for w in words:
        words_length.append(len(w))
    for t in sorted(zip(words_length, words)):
        sorted_words.append(t[1])
    result_set = defaultdict(set)
    root = TrieTreeNode()
    for w in sorted_words:
        root.insert(w, result_set)
    match_set = set()
    for k,v in result_set.items():
        for i in v:
            match_set.add(i)
    unmatch_set = set(words) - match_set
    print unmatch_set | set(result_set.keys())

Answer 1

The best solution I have for this. If my understanding is good, this is a fairly simple to solve and you do not need anything else than a set.

There might be a better way than rebuilding entirely the set.

words = ['foo', 'foog', 'food', 'asdf']
prefixes = set()

for w in words:
    if any(p for p in prefixes if w.startswith(p)):
        # A prefix exists in the prefixes set.
        continue

    # Re-build the set with all elements not matching the set.
    prefixes = set(p for p in prefixes if not p.startswith(w))
    prefixes.add(w)

print(prefixes)

Answer 2

Follow PEP8, and upgrade to Python 3.

---

If you want to keep the Trie, then you should simplify your code. You can get the code to work with insert, without the if or result_set. You also don't need the self.count, and so I'd change your code to:

class TrieTreeNode:
    def __init__(self):
        self.children = defaultdict(TrieTreeNode)
        self.is_end = False

    def insert(self, word):
        node = self
        for char in word:
            node = node.children[char]
        node.isEnd = True


if __name__ == "__main__":
    words = ['foo', 'foog', 'food', 'asdf']
    root = TrieTreeNode()
    for word in words:
        root.insert(word)

After this you just need a loop that will tell you the roots. The children from root all have one child, otherwise that's the end of the word, or if it reaches the end of a word. And so you should be able to get:

output = []
for char, child in root.children.items():
    word = [char]
    while not child.is_end and len(child.children) == 1:
        children = child.children
        letter = children.keys()[0]
        word.append(letter)
        child = children[letter]
    output.append(''.join(word))
print output