main.py (cf8eb64) - dost

main.py

import itertools
import math
import re
import sys
import time
import turtle
import typing
import collections.abc
import Levenshtein

Func: typing.TypeAlias = collections.abc.Callable
DistTable: typing.TypeAlias = list[list[float]]
IndSeq: typing.TypeAlias = list[int]

def tree_weight(x) -> int:
    if x is None:
        return 0
    elif isinstance(x, BinaryTree):
        return x.weight
    return 1

class BinaryTree:
    def __init__(self, a, b=None):
        self.x = self.y = self.parent = None
        self.a = a
        self.b = b
        self.weight = tree_weight(a) + tree_weight(b)
        self.usage = 0

def __str__(self) -> str:
        return str(self.a) if self.b is None else f"({self.a} {self.b})"

def merge(self, other):
        a, b = self, other
        while a.parent:
            a = a.parent
        while b.parent:
            b = b.parent
        if a is b:
            return a
        parent = BinaryTree(a, b)
        a.parent = b.parent = parent
        return parent

def use(self):
        assert self.usage < self.weight
        self.usage += 1
        if self.parent:
            self.parent.use()

def distance_table[T](points: list[T], metric: Func[[T, T], float]) -> DistTable:
    return [[metric(x, y) for y in points] for x in points]

def show_mat(distances: DistTable) -> str:
    rows = []
    for i in range(len(distances)):
        first, last = i == 0, i == len(distances) - 1
        row = "/" if first else "\\" if last else "|"
        row += " ".join(f"{d:5.2f}" for d in distances[i])
        row += "\\" if first else "/" if last else "|"
        rows.append(row)
    return "\n".join(rows)

# this and other quirks in cachey_nearest are legitimately for speed
def min2(a, b):
    return a if a < b else b

# good results, slow: O(n^3)
def cachey_nearest(distances: DistTable) -> IndSeq:
    seq = []
    options = set(range(len(distances)))
    nearests = [math.inf for _ in range(len(distances))]
    while options:
        best, brds, score = None, None, None
        for o in options:
            reduced = [min2(x, y) for x, y in zip(distances[o], nearests)]
            sr = sum(reduced)
            if best is None or sr < score:
                best, brds, score = o, reduced, sr
        nearests = brds
        seq.append(best)
        options.remove(best)
    return seq

# good results, sketchy on a couple trials
def cachey_maximin(distances: DistTable) -> IndSeq:
    seq = []
    options = set(range(len(distances)))
    start = min(options, key=lambda o: sum(distances[o]))
    seq.append(start)
    options.remove(start)
    score = math.inf
    while options:
        updateds = {o: min2(min(distances[o][i] for i in seq), score) for o in options}
        best = max(options, key=lambda o: updateds[o])
        score = updateds[best]
        seq.append(best)
        options.remove(best)
    return seq

def distance_hierarchy(distances: DistTable) -> BinaryTree:
    n = len(distances)
    forest = [BinaryTree(i) for i in range(n)]
    p = None
    for (i, j) in sorted(
        ((i, j) for i in range(n) for j in range(n)),
        key=lambda ij: distances[ij[0]][ij[1]]
    ):
        p = forest[i].merge(forest[j])
    return p

# good results, sketchy on somewhat more trials
def scattered_hierarchy(hierarchy: BinaryTree) -> IndSeq:
    seq = []
    while hierarchy.usage < hierarchy.weight:
        fb = hierarchy
        while fb.b is not None:
            ua = fb.a.usage / fb.a.weight
            ub = fb.b.usage / fb.b.weight
            if ua < ub:
                fb = fb.a
            elif ua == ub and fb.a.weight > fb.b.weight:
                fb = fb.a
            else:
                fb = fb.b
        fb.use()
        seq.append(fb.a)
    return seq

def total_len(distances: DistTable, seq: IndSeq) -> float:
    return sum(distances[i][j] for (i, j) in itertools.pairwise(seq))

points: list[str] = [l.strip() for l in sys.stdin]
to = time.time()
dt: DistTable = distance_table(points, lambda a, b: Levenshtein.distance(a, b, weights=(1, 1, 3)))
tf = time.time()
print(f"took {tf - to} s to calculate distances")
to = time.time()
h: BinaryTree = distance_hierarchy(dt)
tf = time.time()
print(f"took {tf - to} s to build hierarchy")
print(h)
METHODS = [
    ("graph", lambda: scattered_hierarchy(h)),
    ("nearest-nb", lambda: cachey_nearest(dt)),
    ("maximin", lambda: cachey_maximin(dt)),
]
for (name, method) in METHODS:
    to = time.time()
    seq = method()
    tf = time.time()
    l = total_len(dt, seq)
    print(f"{name} method took {tf - to} s to find {seq}, length {l}")
    for k in seq[:int(math.sqrt(len(points)))]:
        print("\t" + points[k])

dost

Frantically optimise cachey_nearest