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.
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.
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: