5
\$\begingroup\$

Today I present some code for finding all words that begin with a given prefix and some simple test cases for review.

from itertools import chain


sentinel = object()


def insert(tree, word):
    for letter in chain(word, [sentinel]):
        tree = tree.setdefault(letter, {})


def words_beginning_with(start, tree):
    for letter in start:
        try:
            tree = tree[letter]
        except KeyError:
            return
    for word in _all_words(tree):
        yield start + word


def _all_words(tree, start=''):
    for letter, tree in tree.items():
        if letter is sentinel:
            yield start
        else:
            for word in _all_words(tree, start+letter):
                yield word


def test_tree():
    # Simple tests
    tree = {}

    assert [] == list(words_beginning_with('foo', tree))

    insert(tree, 'foo')
    assert {'f': {'o': {'o': {sentinel: {}}}}} == tree

    assert ['foo'] == list(words_beginning_with('foo', tree))

    insert(tree, 'food')
    assert {'f': {'o': {'o': {sentinel: {},
                              'd': {sentinel: {}}}}}} == tree

    assert sorted(['foo', 'food']) == sorted(words_beginning_with('fo', tree))

    insert(tree, 'foody')
    assert sorted(['foo', 'food', 'foody']) == sorted(
        words_beginning_with('fo', tree))

    insert(tree, 'fold')
    assert sorted(['foo', 'fold', 'food', 'foody']) == sorted(
        words_beginning_with('fo', tree))


if __name__ == "__main__":
    test_tree()

I appreciate brutal and honest feedback the most, thanks.

\$\endgroup\$
  • \$\begingroup\$ Any reason not to take advantage of startwith() ? \$\endgroup\$ – boardrider Dec 2 '16 at 14:15
  • \$\begingroup\$ If you mean just iterating over all words added and checking it with string.startswith I imagine it would be a lot slower if the list of words is big (haven't benchmarked). If you mean making the structure of the tree smaller by converting it to a radix tree, possibly. I was going to convert to a radix tree as an answer to this post. \$\endgroup\$ – Justin Fay Dec 3 '16 at 13:52
4
\$\begingroup\$
  • Sentinel could simply be None instead of a "custom" object.
  • You don't need chain either: just set the sentinel manually at the end of the for loop.
  • If using Python 3, you can use yield from ... instead of from ...: yield.
def insert(tree, word):
    for letter in word:
        tree = tree.setdefault(letter, {})
    tree[None] = {}  # or = None, it doesn't matter


def words_beginning_with(start, tree):
    for letter in start:
        try:
            tree = tree[letter]
        except KeyError:
            return
    for word in _all_words(tree):
        yield start + word


def _all_words(tree, start=''):
    for letter, tree in tree.items():
        if letter is None:
            yield start
        else:
            yield from _all_words(tree, start + letter):
\$\endgroup\$
  • \$\begingroup\$ I'm curious, do you not think a class would improve the code? \$\endgroup\$ – Peilonrayz Dec 1 '16 at 16:19
  • \$\begingroup\$ @Peilonrayz For such simple task/use case, not necessarily. But if you think otherwise, feel free to suggest it yourself. ;) \$\endgroup\$ – Mathias Ettinger Dec 1 '16 at 16:26
  • \$\begingroup\$ Thanks. It was just simple curiosity. You've stated all my other points anyway, :P \$\endgroup\$ – Peilonrayz Dec 1 '16 at 16:32

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.