import math
import typing
import collections.abc

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 score_nearest(distances: DistTable, sample: IndSeq) -> float:
    return sum(min(distances[i][j] for j in sample) for i in range(len(distances)))

def score_total(distances: DistTable, sample: IndSeq) -> float:
    return sum(sum(
        0 if i in sample else distances[i][j] for j in sample
    ) for i in range(len(distances)))

def greedy_min_seq(distances: DistTable, score_func: Func[[DistTable, IndSeq], float]) -> IndSeq:
    seq = []
    options = set(range(len(distances)))
    while options:
        best = min(options, key=lambda o: score_func(distances, seq + [o]))
        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

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

dt: DistTable = distance_table([(0, 2), (3, 2), (4, 2), (5, 0), (0, 0)], math.dist)
print(greedy_min_seq(dt, score_total))
h: BinaryTree = distance_hierarchy(dt)
print(h)
print(scattered_hierarchy(h))