File size: 2,686 Bytes
58258b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Shared primitives: union-find, cluster -> representative selection, normalize."""
from __future__ import annotations
import re
from collections import defaultdict
from typing import Any, Callable, Iterable
class UnionFind:
"""Union-find with path compression and union-by-rank."""
def __init__(self, n: int):
self.parent = list(range(n))
self.rank = [0] * n
def find(self, x: int) -> int:
root = x
while self.parent[root] != root:
root = self.parent[root]
while self.parent[x] != root:
self.parent[x], x = root, self.parent[x]
return root
def union(self, x: int, y: int) -> None:
rx, ry = self.find(x), self.find(y)
if rx == ry:
return
if self.rank[rx] < self.rank[ry]:
rx, ry = ry, rx
self.parent[ry] = rx
if self.rank[rx] == self.rank[ry]:
self.rank[rx] += 1
def cluster_and_pick(
n: int,
pairs: Iterable[tuple[int, int]],
key_fn: Callable[[int], Any] | None = None,
) -> tuple[list[int], dict[int, list[int]]]:
"""Cluster items in [0, n) via `pairs`, pick one rep per cluster.
Returns (keep, clusters):
keep: sorted list of representative indices (one per cluster).
clusters: dict {rep_idx -> sorted list of ALL member indices in that
cluster, including the rep}. Singletons appear as {i: [i]}.
"""
uf = UnionFind(n)
for i, j in pairs:
uf.union(i, j)
members_by_root: dict[int, list[int]] = defaultdict(list)
for i in range(n):
members_by_root[uf.find(i)].append(i)
clusters: dict[int, list[int]] = {}
for members in members_by_root.values():
rep = min(members) if key_fn is None else min(members, key=key_fn)
clusters[rep] = sorted(members)
keep = sorted(clusters.keys())
return keep, clusters
def pick_representatives(
n: int,
pairs: Iterable[tuple[int, int]],
key_fn: Callable[[int], Any] | None = None,
) -> list[int]:
"""Cluster items in [0, n) via `pairs`, then keep one rep per cluster.
Representative is the cluster member with the smallest `key_fn(i)`. With
no key_fn, falls back to smallest index ("first seen"). Returned indices
are sorted ascending so caller can re-index in original order.
Thin wrapper over `cluster_and_pick` that drops the cluster-membership map.
"""
keep, _ = cluster_and_pick(n, pairs, key_fn)
return keep
_WHITESPACE_RE = re.compile(r"\s+")
def normalize_text(text: str) -> str:
"""Lowercase + collapse internal whitespace + strip."""
return _WHITESPACE_RE.sub(" ", text.lower()).strip()
|