Working solution for M

autocomplete.py

@@ -0,0 +1,140 @@
+#!/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)))
+

autocomplete.txt

@@ -0,0 +1,30 @@
+5
+5
+hi
+hello
+lol
+hills
+hill
+5
+a
+aa
+aaa
+aaaa
+aaaaa
+5
+aaaaa
+aaaa
+aaa
+aa
+a
+6
+to
+be
+or
+not
+two
+bee
+3
+having
+fun
+yet