Search iteratively in a ternary search tree

Question

I was trying to write an iterative search function for my Ternary Search Tree class based on a pseudo-code, and now it seems to work, but I think it can definitely be improved.

def search(self, key: str):
    return TST._search(self.root, key)

@staticmethod
def _search(node: TSTNode, key: str):
    """Note that node.key is a character"""
    if node is None or not key:
        return None

    for i in range(len(key) - 1):

        while node is not None and key[i] != node.key:
            if key[i] < node.key:
                node = node.left
            else:
                node = node.right

        if node is None:  # Unsuccessful search
            return None
        else:
            node = node.mid

    # In case the length of the key is 1 (but NOT only, see edit below)
    return node.value if node.key == key else None

Is my algorithm correct? My algorithm does not allow keys or characters to be empty strings.

---

Edit 1

I think there's a small bug in the code above in the last statement, which should be:

if not node or node.key != key[-1]:
    return None
else:
    return node.value

This is because there's one more way to arrive at that return statement (apart from having initially the key of length 1, as I wrote in the comment), which is to exit from the while loop because the condition key[i] != node.key is false.

In that case, we would be comparing a character, i.e. node.key, with a string key (if the node is not None: see note below!). In Python this can be easily hidden leading to bugs like this one.

Note also that node could still be None (I think) after we do node = node.mid.

Edit 2

After testing this method, I noticed that the last correction of the previous edit didn't solve all problems with the function. The last problem is that we still may need to go left or right after we exit the for loop. Here you can find the correct version:

https://github.com/nbro/ands/blob/master/ands/ds/TST.py

Answer 1

Iterables

Generally speaking, if you use range(len(some_iterable))and then access an index of the item, you can just iterate over the item itself:

for i in range(len(word)):  # okay
    print(word[i])

for character in word:  # better
    print(character)

In your case it'd be key[:-1]

Style

Good on you for following PEP8. The code looks beautiful. There are more blank lines than I usually see, but they work.

The ternary operator (x if y else z)

Your use of the ternary operator is quite elegant here, but be aware that this makes your code harder to test.

I often see code that use the ternary operator as a sort of cheat to get out of unit test coverage (coverage.py doesn't notice if only one of the cases in a ternary statement is executed).

The algorithm

It looks fine to me. I'm not certain about the correctness. Can you link that pseudo-code you're working from?

You can make it a bit more obvious if you like. How do I call this method? What kind of object is key? I'm sure this makes complete sense in context with the rest of the class, but you can always add argument annotations and function annotations to make your code that much more readable:

def nice_number(x: float) -> str:
    ...

or

def f(n: 'must be prime'):
    ...

What can you do to improve this?

Write a unit test. It makes it so much easier to play with the implementation, to try out different things, without having to worry about breaking some edge case. Write a complete unit test, and be aware of what the algorithm can an can't do.

Run profiling. I don't think you will learn anything shocking by profiling the code, but it's always good to get a sense of what takes the most amount of time. My guess is that the slowest part of your code, considering how straightforward it is, will be the attribute access (node.key). (Rule of thumb: dots are expensive)

Write a docstring. Run help() on your module and try to imagine what it looks like to users.

Get rid of the wrapper method. Unless there's something I'm not seeing here, there's no reason to have a static method TST._search(), not visible to help() because of that underscore, and a wrapper method TST.search() that merely calls it.

If you do get rid of the static method, start off by setting node = self.root (going back to the earlier comment about dots being expensive, you don't want to have that in the hot part of your loop if avoiding it is so easy).

I almost never use staticmethod() in real-world code. If I do, I make sure to spend a minute thinking about whether it's really necessary.