ProgrammingChallenges/autocomplete.py

#!/usr/bin/env python3

# Tree data structure
class Node:
    def __init__(self, key):
        self.key = key
        self.cost = None

        self.left = None
        self.right = None
        self.parent = None

    def set_left(self, left):
        self.left = left
        left.parent = self

    def set_right(self, right):
        self.right = right
        right.parent = self

    def string(self):
        return "(" + ("" if self.left is None else Node.string(self.left) + ",") + \
                self.key + ("" if self.right is None else "," + Node.string(self.right)) + ")"


def tree_insert(tree, key):
    if tree.key == key:
        return (tree, True)
    elif key < tree.key:
        if tree.left is None:
            tree.set_left(Node(key))
            return (tree.left, False)
        else:
            return tree_insert(tree.left, key)
    else:
        if tree.right is None:
            tree.set_right(Node(key))
            return (tree.right, False)
        else:
            return tree_insert(tree.right, key)


def tree_min(t):
    if t is None:
        return None
    while t.left is not None:
        t = t.left
    return t


def tree_max(t):
    if t is None:
        return None
    while t.right is not None:
        t = t.right
    return t


def tree_successor(t):
    if t.right is not None:
        return tree_min(t.right)
    while t.parent is not None:
        if t.parent.left == t:
            return t.parent
        else:
            t = t.parent
    return None


def tree_predecessor(t):
    if t.left is not None:
        return tree_max(t.left)
    while t.parent is not None:
        if t.parent.right == t:
            return t.parent
        else:
            t = t.parent
    return None


def compute_cost(reference, word):
    i = 0
    while i < len(reference) and i < len(word):
        if word[i] != reference[i]:
            return i + 1
        i += 1
    if i == len(reference) and i < len(word):
        return i + 1
    else:
        return i


def autocomplete(W):
    tree = Node(W[0])
    tree.cost = 1
    count = tree.cost
    # print(tree.key, tree.cost)

    for i in range(1, len(W)):
        node, already_there = tree_insert(tree, W[i])
        if already_there:
            count += node.cost
        else:
            prec = tree_predecessor(node)
            succ = tree_successor(node)

            # print("prec:", None if prec is None else prec.key, \
            #         "succ:", None if succ is None else succ.key)

            prec_cost = None
            succ_cost = None
            if prec is not None:
                prec_cost = compute_cost(prec.key, node.key)
            if succ is not None:
                succ_cost = compute_cost(succ.key, node.key)

            # print("prec_cost:", prec_cost, "succ_cost:", succ_cost)

            if prec_cost is not None and succ_cost is not None:
                node.cost = max(prec_cost, succ_cost)
            elif prec_cost is not None:
                node.cost = prec_cost
            elif succ_cost is not None:
                node.cost = succ_cost

            count += node.cost
        # print(tree.string())
        # print(node.key, node.cost)
    return count


if __name__ == "__main__":
    t = int(input())
    for i in range(1, t + 1):
        n = int(input())
        W = []
        for j in range(n):
            W.append(input())
        print("Case #" + str(i) + ": " + str(autocomplete(W)))