apcamargo · August 3, 2026 12:23
diff --git a/partition_align.py b/partition_align.py
 """
 Find correspondences between groups belonging to different partitions of the
 same fixed collection of elements.

 Design choices implemented here:
  * strict one-to-one correspondences (no splits / merges)
  * deterministic tie-breaking -> one canonical answer, always
  * multiple partitions are aligned against a single reference partition
  * optional weighting, at two independent levels (see below)

 Complexity
 ----------
  n = number of elements, K = max number of groups in a partition, m = partitions

  contingency table   O(n)          (exploits: each element sits in one group)
  similarity matrix   O(#nonzero)   (<= n entries, not K^2)
  assignment          O(K^3)        (Hungarian. K is usually << n)
  full alignment      O(m * (n + K^3))

 Weighting is free: it changes the numbers in the cost matrix, not the algorithm.
 The one requirement is *separability* -- a pair's score must not depend on which
 other pairs were chosen.  Both weighting schemes below respect that, so this
 remains a linear assignment problem solved exactly by the Hungarian algorithm.

 The two weighting levels
 ------------------------
 1. ELEMENT WEIGHTS (`element_weights`) -- "these elements matter more".
   |A n B| becomes the summed weight of the shared elements.  Similarity stays
   in [0, 1] and the edit-distance equivalence survives exactly, because the
   weighted symmetric difference telescopes just like the unweighted one.
   Group importance emerges naturally as the sum of its members' weights.

 2. GROUP PRIORITY (`priorities`) -- "groups of this category win ties and may
   override".  A multiplier on the *normalised* similarity:

       score(i, j) = similarity(i, j) * priority(i, j)

   Do NOT fold group size into this multiplier: similarity is already
   size-normalised, and re-introducing size lets a large, poorly-matching group
   outbid a small, near-perfect one.  With a pure multiplier the semantics are
   exact -- weight 2 means "willing to accept a match half as good".

 Both default to 1.0, and with the defaults the output is bit-identical to the
 unweighted version.

 Why no explicit edit distance
 -----------------------------
 The set edit distance between two groups is the symmetric difference
 |P d Q| = |P| + |Q| - 2|P n Q|.  Summed over a complete matching this telescopes
 to `2n - 2 * S|P n Q|`, so *minimising total edit distance is exactly equivalent
 to maximising total overlap*.  Distances never need to be materialised.
 """

 import math
 from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Callable, Dict, Hashable, Iterable, List, Sequence, Tuple

 # A partition is a sequence of groups. A group is any iterable of elements.
 Partition = Sequence[Iterable[Hashable]]
 GroupId = Tuple[int, int]  # (partition_index, group_index)
 Priorities = Sequence[float]  # one multiplier per group, parallel to a partition


 # --------------------------------------------------------------------------- #
 # 1. Contingency table (optionally element-weighted)                          #
 # --------------------------------------------------------------------------- #


 def contingency(
    a: Partition, b: Partition, element_weights: Dict[Hashable, float] | None = None
 ) -> Tuple[Dict[Tuple[int, int], float], List[float], List[float]]:
    """Weighted intersection sizes for every *overlapping* pair, in O(n).

    With `element_weights=None` every element counts 1.0 and this is the plain
    contingency table.  Otherwise each element contributes its own weight.

    Returns (counts, sizes_a, sizes_b).  `counts` is sparse: at most n entries,
    so the bipartite graph stays sparse even with many groups.
    """
    if element_weights:

        def w(e: Hashable) -> float:
            return element_weights.get(e, 1.0)
    else:

        def w(e: Hashable) -> float:
            return 1.0

    sizes_a = [float(sum(w(e) for e in set(g))) for g in a]
    sizes_b = [float(sum(w(e) for e in set(g))) for g in b]

    # element -> indices of the groups containing it (a list, so that partitions
    # with overlapping groups still work. For true partitions it has length 1)
    where_b: Dict[Hashable, List[int]] = defaultdict(list)
    for j, group in enumerate(b):
        for e in set(group):
            where_b[e].append(j)

    counts: Dict[Tuple[int, int], float] = defaultdict(float)
    for i, group in enumerate(a):
        for e in set(group):
            hits = where_b.get(e)
            if hits:
                we = w(e)
                for j in hits:
                    counts[(i, j)] += we
    return dict(counts), sizes_a, sizes_b


 # --------------------------------------------------------------------------- #
 # 2. Similarity                                                               #
 # --------------------------------------------------------------------------- #


 def jaccard(c: float, size_a: float, size_b: float) -> float:
    union = size_a + size_b - c
    return c / union if union > 0 else 0.0


 def dice(c: float, size_a: float, size_b: float) -> float:
    total = size_a + size_b
    return 2 * c / total if total > 0 else 0.0


 def overlap_count(c: float, size_a: float, size_b: float) -> float:
    """Raw |a n b|: the exact edit-distance objective biased toward big groups."""
    return c


 METRICS: Dict[str, Callable[[float, float, float], float]] = {
    "jaccard": jaccard,
    "dice": dice,
    "overlap": overlap_count,
 }


 # --------------------------------------------------------------------------- #
 # 3. Priority combination                                                     #
 # --------------------------------------------------------------------------- #
 #
 # The two groups in a pair may carry different priorities, so a combination rule
 # is needed.  `max` is the usual choice for priority semantics ("either group
 # being important is enough"). `min` demands both. `geometric` compromises.
 # `reference` ignores the second partition entirely, which is handy when priority
 # is a property of the reference partition's categories only.

 COMBINERS: Dict[str, Callable[[float, float], float]] = {
    "max": max,
    "min": min,
    "geometric": lambda x, y: math.sqrt(x * y),
    "reference": lambda x, y: x,
 }


 def _priority_of(
    pa: Priorities | None,
    pb: Priorities | None,
    i: int,
    j: int,
    combine: Callable[[float, float], float],
 ) -> float:
    x = pa[i] if pa is not None else 1.0
    y = pb[j] if pb is not None else 1.0
    if x <= 0 or y <= 0:
        raise ValueError("priorities must be strictly positive")
    if pa is None and pb is None:
        return 1.0
    return combine(x, y)


 # --------------------------------------------------------------------------- #
 # 4. Deterministic scoring                                                    #
 # --------------------------------------------------------------------------- #
 #
 # A max-weight matching can have several optima with identical total weight.
 # To always return the same canonical answer we pack a lexicographic key
 #     (weighted similarity, shared weight, prefer-low-indices)
 # into one big integer.  The bases are chosen so that the low-order components
 # can never sum high enough to disturb a higher-order one, hence maximising the
 # packed sum == lexicographic optimisation of the three criteria in order.
 # Python's arbitrary-precision ints make this exact (no float drift).

 _SIM_SCALE = 10**9  # resolution of the primary score
 _SEC_SCALE = 10**6  # resolution of the secondary (shared-weight) score


 def _pack_scores(
    counts,
    sizes_a,
    sizes_b,
    metric,
    threshold,
    weight_total,
    priority_a=None,
    priority_b=None,
    combine=max,
 ):
    """Sparse dicts {(i, j): packed_int} and {(i, j): raw_similarity}.

    `threshold` is applied to the RAW similarity, before the priority
    multiplier -- otherwise a high-priority pair could sneak a structurally
    poor match past the cutoff.
    """
    ka, kb = len(sizes_a), len(sizes_b)
    k_max = max(ka, kb, 1)

    tert_bound = k_max * k_max + k_max + 1  # tertiary < tert_bound
    base_sec = k_max * tert_bound  # blocks tertiary carry
    sec_bound = int(weight_total * _SEC_SCALE) + 1  # secondary < sec_bound
    base_pri = base_sec * (k_max * sec_bound + 1)  # blocks secondary carry

    packed: Dict[Tuple[int, int], int] = {}
    sims: Dict[Tuple[int, int], float] = {}
    for (i, j), c in counts.items():
        raw = metric(c, sizes_a[i], sizes_b[j])
        if raw < threshold or raw <= 0.0:
            continue
        effective = raw * _priority_of(priority_a, priority_b, i, j, combine)
        primary = int(round(effective * _SIM_SCALE))
        secondary = min(int(round(c * _SEC_SCALE)), sec_bound - 1)  # shared weight
        tertiary = (k_max - i) * k_max + (k_max - j)  # lower ids win
        packed[(i, j)] = primary * base_pri + secondary * base_sec + tertiary
        sims[(i, j)] = raw
    return packed, sims


 # --------------------------------------------------------------------------- #
 # 5. Assignment                                                               #
 # --------------------------------------------------------------------------- #


 def _hungarian(cost: List[List[int]]) -> List[int]:
    """Min-cost assignment on a square integer matrix (Jonker-Volgenant style
    successive shortest paths).  Returns row -> column.

    Strictly integral: the entries are arbitrary-precision packed scores that
    routinely exceed 2**53, so a float sentinel would risk silently collapsing
    distinct scores.  `big` is derived from the matrix and provably exceeds any
    reduced cost -- potentials grow by at most one delta per iteration, so
    |u| + |v| <= 2k * max_abs, and cur = cost - u - v < (2k + 1) * max_abs.
    """
    k = len(cost)
    if k == 0:
        return []
    max_abs = max((abs(x) for row in cost for x in row), default=0)
    big: int = (4 * k + 8) * (max_abs + 1)

    u: List[int] = [0] * (k + 1)
    v: List[int] = [0] * (k + 1)
    p: List[int] = [0] * (k + 1)
    way: List[int] = [0] * (k + 1)

    for i in range(1, k + 1):
        p[0] = i
        j0 = 0
        minv: List[int] = [big] * (k + 1)
        used: List[bool] = [False] * (k + 1)
        while True:
            used[j0] = True
            i0 = p[j0]
            delta: int = big
            j1 = 0
            for j in range(1, k + 1):
                if not used[j]:
                    cur = cost[i0 - 1][j - 1] - u[i0] - v[j]
                    if cur < minv[j]:
                        minv[j], way[j] = cur, j0
                    if minv[j] < delta:
                        delta, j1 = minv[j], j
            for j in range(k + 1):
                if used[j]:
                    u[p[j]] += delta
                    v[j] -= delta
                else:
                    minv[j] -= delta
            j0 = j1
            if p[j0] == 0:
                break
        while j0:
            j1 = way[j0]
            p[j0] = p[j1]
            j0 = j1

    row_to_col: List[int] = [-1] * k
    for j in range(1, k + 1):
        if p[j]:
            row_to_col[p[j] - 1] = j - 1
    return row_to_col


 def match_two(
    a: Partition,
    b: Partition,
    *,
    metric: str = "jaccard",
    threshold: float = 0.0,
    element_weights: Dict[Hashable, float] | None = None,
    priority_a: Priorities | None = None,
    priority_b: Priorities | None = None,
    combine: str = "max",
    weight_total: float | None = None,
 ) -> Tuple[List[Tuple[int, int, float]], List[int], List[int]]:
    """Optimal, deterministic 1-1 correspondence between two partitions.

    Returns (pairs, unmatched_a, unmatched_b). Pairs are (i, j, raw_similarity).
    The similarity reported is the raw, un-boosted one, so it stays comparable
    across pairs with different priorities.
    """
    fn = METRICS[metric]
    comb = COMBINERS[combine]
    counts, sizes_a, sizes_b = contingency(a, b, element_weights)
    if weight_total is None:
        weight_total = max(sum(sizes_a), sum(sizes_b), 1.0)

    packed, sims = _pack_scores(
        counts,
        sizes_a,
        sizes_b,
        fn,
        threshold,
        weight_total,
        priority_a,
        priority_b,
        comb,
    )

    ka, kb = len(a), len(b)
    k = max(ka, kb)
    if k == 0:
        return [], [], []

    # square, padded. Padding and sub-threshold pairs score 0 -> stay unmatched
    cost = [[0] * k for _ in range(k)]
    for (i, j), score in packed.items():
        cost[i][j] = -score  # maximise weight == minimise negative weight

    row_to_col = _hungarian(cost)

    pairs, matched_a, matched_b = [], set(), set()
    for i in range(ka):
        j = row_to_col[i]
        if j < kb and (i, j) in sims:
            pairs.append((i, j, sims[(i, j)]))
            matched_a.add(i)
            matched_b.add(j)
    pairs.sort(key=lambda t: (-t[2], t[0], t[1]))

    unmatched_a = [i for i in range(ka) if i not in matched_a]
    unmatched_b = [j for j in range(kb) if j not in matched_b]
    return pairs, unmatched_a, unmatched_b


 def greedy_match(
    a: Partition,
    b: Partition,
    *,
    metric: str = "jaccard",
    threshold: float = 0.0,
    element_weights: Dict[Hashable, float] | None = None,
    priority_a: Priorities | None = None,
    priority_b: Priorities | None = None,
    combine: str = "max",
 ):
    """O(n log n) fallback for very large numbers of groups.
    Guaranteed within 1/2 of the optimum. Near-exact on similar partitions."""
    fn = METRICS[metric]
    comb = COMBINERS[combine]
    counts, sizes_a, sizes_b = contingency(a, b, element_weights)
    cand = []
    for (i, j), c in counts.items():
        raw = fn(c, sizes_a[i], sizes_b[j])
        if raw > threshold and raw > 0:
            eff = raw * _priority_of(priority_a, priority_b, i, j, comb)
            cand.append((-eff, -c, i, j, raw))
    cand.sort(key=lambda t: t[:4])
    used_a, used_b, pairs = set(), set(), []
    for _, _, i, j, raw in cand:
        if i not in used_a and j not in used_b:
            used_a.add(i)
            used_b.add(j)
            pairs.append((i, j, raw))
    unmatched_a = [i for i in range(len(sizes_a)) if i not in used_a]
    unmatched_b = [j for j in range(len(sizes_b)) if j not in used_b]
    return pairs, unmatched_a, unmatched_b


 # --------------------------------------------------------------------------- #
 # 6. Multi-partition alignment (reference-based)                              #
 # --------------------------------------------------------------------------- #


 @dataclass
 class Correspondence:
    """One correspondence class: at most one group per partition."""

    members: Dict[int, int] = field(default_factory=dict)  # partition -> group
    similarities: Dict[int, float] = field(default_factory=dict)  # raw, vs reference

    def group_ids(self) -> List[GroupId]:
        return sorted(self.members.items())


 @dataclass
 class Alignment:
    reference: int
    classes: List[Correspondence]

    def render(self, partitions: Sequence[Partition]) -> str:
        lines = [f"reference partition = {self.reference}"]
        for n, cls in enumerate(self.classes):
            parts = []
            for g, idx in sorted(cls.members.items()):
                elems = sorted(map(str, partitions[g][idx]))
                sim = cls.similarities.get(g)
                tag = "" if sim is None else f"~{sim:.2f}"
                parts.append(f"G{g}[{','.join(elems)}]{tag}")
            lines.append(f"  {n}: " + "  <->  ".join(parts))
        return "\n".join(lines)


 def _pick_reference(partitions, metric, element_weights, priorities, combine) -> int:
    """Medoid: the partition most similar to all the others (cheap greedy score)."""
    best, best_score = 0, -1.0
    for r in range(len(partitions)):
        total = 0.0
        for o in range(len(partitions)):
            if o == r:
                continue
            pairs, _, _ = greedy_match(
                partitions[r],
                partitions[o],
                metric=metric,
                element_weights=element_weights,
                priority_a=priorities[r] if priorities else None,
                priority_b=priorities[o] if priorities else None,
                combine=combine,
            )
            total += sum(s for _, _, s in pairs)
        if total > best_score:
            best, best_score = r, total
    return best


 def align(
    partitions: Sequence[Partition],
    *,
    reference: int | str = "medoid",
    metric: str = "jaccard",
    threshold: float = 0.0,
    element_weights: Dict[Hashable, float] | None = None,
    priorities: Sequence[Priorities] | None = None,
    combine: str = "max",
    exact: bool = True,
 ) -> Alignment:
    """Align every partition to one reference partition.

    `element_weights` maps an element to its importance (default 1.0).
    `priorities` is parallel to `partitions`: one multiplier per group.

    Groups with no counterpart in the reference become their own correspondence
    class (like `[B]` in the example).  Note the consequence of the
    reference-based strategy: two non-reference partitions that share an orphan
    group are not linked to each other, because all links go through the
    reference.
    """
    if not partitions:
        return Alignment(0, [])
    if priorities is not None and len(priorities) != len(partitions):
        raise ValueError("`priorities` must have one entry per partition")

    if reference == "medoid":
        ref = _pick_reference(partitions, metric, element_weights, priorities, combine)
    else:
        ref = int(reference)

    # a single shared bound keeps the integer packing consistent across pairings
    weight_total = (
        max(
            (
                sum(
                    element_weights.get(e, 1.0) if element_weights else 1.0
                    for e in {x for grp in g for x in grp}
                )
                for g in partitions
            ),
            default=1.0,
        )
        or 1.0
    )

    classes = [Correspondence(members={ref: i}) for i in range(len(partitions[ref]))]
    matcher = match_two if exact else greedy_match

    for g in range(len(partitions)):
        if g == ref:
            continue
        kwargs = dict(
            metric=metric,
            threshold=threshold,
            element_weights=element_weights,
            priority_a=priorities[ref] if priorities else None,
            priority_b=priorities[g] if priorities else None,
            combine=combine,
        )
        if exact:
            kwargs["weight_total"] = weight_total
        pairs, _, unmatched_other = matcher(partitions[ref], partitions[g], **kwargs)
        for i, j, sim in pairs:
            classes[i].members[g] = j
            classes[i].similarities[g] = sim
        for j in unmatched_other:  # orphan -> own class
            classes.append(Correspondence(members={g: j}, similarities={g: 0.0}))

    return Alignment(ref, classes)


 # --------------------------------------------------------------------------- #
 # Demo                                                                        #
 # --------------------------------------------------------------------------- #

 if __name__ == "__main__":
    partitions = [
        [["A", "B", "C", "D"], ["E", "F"]],  # partition 0
        [["A", "C", "D"], ["E", "F"], ["B"]],  # partition 1
        [["A", "C"], ["D", "B"], ["E", "F"]],  # partition 2
    ]

    print("--- unweighted (unchanged from before) ---")
    print(align(partitions[:2], reference=0).render(partitions[:2]))

    print("\n--- element weights: B is important, so [B] matters more ---")
    print(
        align(partitions[:2], reference=0, element_weights={"B": 10.0}).render(
            partitions[:2]
        )
    )

    # two groups in `right` compete for the same partner in `left`
    left = [["A", "B", "C"]]
    right = [["A", "B", "C", "D"], ["A"]]
    print("\n--- group priority overriding a better structural match ---")
    for pri in (None, [[1.0], [1.0, 3.0]]):
        pairs, _, _ = match_two(
            left,
            right,
            priority_a=pri[0] if pri else None,
            priority_b=pri[1] if pri else None,
        )
        i, j, s = pairs[0]
        print(
            f"  priorities={pri}: -> right[{j}] "
            f"{sorted(right[j])} at raw similarity {s:.2f}"
        )
	"""
	Find correspondences between groups belonging to different partitions of the
	same fixed collection of elements.

	Design choices implemented here:
	* strict one-to-one correspondences (no splits / merges)
	* deterministic tie-breaking -> one canonical answer, always
	* multiple partitions are aligned against a single reference partition
	* optional weighting, at two independent levels (see below)

	Complexity
	----------
	n = number of elements, K = max number of groups in a partition, m = partitions

	contingency table O(n) (exploits: each element sits in one group)
	similarity matrix O(#nonzero) (<= n entries, not K^2)
	assignment O(K^3) (Hungarian. K is usually << n)
	full alignment O(m * (n + K^3))

	Weighting is free: it changes the numbers in the cost matrix, not the algorithm.
	The one requirement is separability -- a pair's score must not depend on which
	other pairs were chosen. Both weighting schemes below respect that, so this
	remains a linear assignment problem solved exactly by the Hungarian algorithm.

	The two weighting levels
	------------------------
	1. ELEMENT WEIGHTS (`element_weights`) -- "these elements matter more".
	\|A n B\| becomes the summed weight of the shared elements. Similarity stays
	in [0, 1] and the edit-distance equivalence survives exactly, because the
	weighted symmetric difference telescopes just like the unweighted one.
	Group importance emerges naturally as the sum of its members' weights.

	2. GROUP PRIORITY (`priorities`) -- "groups of this category win ties and may
	override". A multiplier on the normalised similarity:

	score(i, j) = similarity(i, j) * priority(i, j)

	Do NOT fold group size into this multiplier: similarity is already
	size-normalised, and re-introducing size lets a large, poorly-matching group
	outbid a small, near-perfect one. With a pure multiplier the semantics are
	exact -- weight 2 means "willing to accept a match half as good".

	Both default to 1.0, and with the defaults the output is bit-identical to the
	unweighted version.

	Why no explicit edit distance
	-----------------------------
	The set edit distance between two groups is the symmetric difference
	\|P d Q\| = \|P\| + \|Q\| - 2\|P n Q\|. Summed over a complete matching this telescopes
	to `2n - 2 * S\|P n Q\|`, so *minimising total edit distance is exactly equivalent
	to maximising total overlap*. Distances never need to be materialised.
	"""

	import math
	from collections import defaultdict
	from dataclasses import dataclass, field
	from typing import Callable, Dict, Hashable, Iterable, List, Sequence, Tuple

	# A partition is a sequence of groups. A group is any iterable of elements.
	Partition = Sequence[Iterable[Hashable]]
	GroupId = Tuple[int, int] # (partition_index, group_index)
	Priorities = Sequence[float] # one multiplier per group, parallel to a partition


	# --------------------------------------------------------------------------- #
	# 1. Contingency table (optionally element-weighted) #
	# --------------------------------------------------------------------------- #


	def contingency(
	a: Partition, b: Partition, element_weights: Dict[Hashable, float] \| None = None
	) -> Tuple[Dict[Tuple[int, int], float], List[float], List[float]]:
	"""Weighted intersection sizes for every overlapping pair, in O(n).

	With `element_weights=None` every element counts 1.0 and this is the plain
	contingency table. Otherwise each element contributes its own weight.

	Returns (counts, sizes_a, sizes_b). `counts` is sparse: at most n entries,
	so the bipartite graph stays sparse even with many groups.
	"""
	if element_weights:

	def w(e: Hashable) -> float:
	return element_weights.get(e, 1.0)
	else:

	def w(e: Hashable) -> float:
	return 1.0

	sizes_a = [float(sum(w(e) for e in set(g))) for g in a]
	sizes_b = [float(sum(w(e) for e in set(g))) for g in b]

	# element -> indices of the groups containing it (a list, so that partitions
	# with overlapping groups still work. For true partitions it has length 1)
	where_b: Dict[Hashable, List[int]] = defaultdict(list)
	for j, group in enumerate(b):
	for e in set(group):
	where_b[e].append(j)

	counts: Dict[Tuple[int, int], float] = defaultdict(float)
	for i, group in enumerate(a):
	for e in set(group):
	hits = where_b.get(e)
	if hits:
	we = w(e)
	for j in hits:
	counts[(i, j)] += we
	return dict(counts), sizes_a, sizes_b


	# --------------------------------------------------------------------------- #
	# 2. Similarity #
	# --------------------------------------------------------------------------- #


	def jaccard(c: float, size_a: float, size_b: float) -> float:
	union = size_a + size_b - c
	return c / union if union > 0 else 0.0


	def dice(c: float, size_a: float, size_b: float) -> float:
	total = size_a + size_b
	return 2 * c / total if total > 0 else 0.0


	def overlap_count(c: float, size_a: float, size_b: float) -> float:
	"""Raw \|a n b\|: the exact edit-distance objective biased toward big groups."""
	return c


	METRICS: Dict[str, Callable[[float, float, float], float]] = {
	"jaccard": jaccard,
	"dice": dice,
	"overlap": overlap_count,
	}


	# --------------------------------------------------------------------------- #
	# 3. Priority combination #
	# --------------------------------------------------------------------------- #
	#
	# The two groups in a pair may carry different priorities, so a combination rule
	# is needed. `max` is the usual choice for priority semantics ("either group
	# being important is enough"). `min` demands both. `geometric` compromises.
	# `reference` ignores the second partition entirely, which is handy when priority
	# is a property of the reference partition's categories only.

	COMBINERS: Dict[str, Callable[[float, float], float]] = {
	"max": max,
	"min": min,
	"geometric": lambda x, y: math.sqrt(x * y),
	"reference": lambda x, y: x,
	}


	def _priority_of(
	pa: Priorities \| None,
	pb: Priorities \| None,
	i: int,
	j: int,
	combine: Callable[[float, float], float],
	) -> float:
	x = pa[i] if pa is not None else 1.0
	y = pb[j] if pb is not None else 1.0
	if x <= 0 or y <= 0:
	raise ValueError("priorities must be strictly positive")
	if pa is None and pb is None:
	return 1.0
	return combine(x, y)


	# --------------------------------------------------------------------------- #
	# 4. Deterministic scoring #
	# --------------------------------------------------------------------------- #
	#
	# A max-weight matching can have several optima with identical total weight.
	# To always return the same canonical answer we pack a lexicographic key
	# (weighted similarity, shared weight, prefer-low-indices)
	# into one big integer. The bases are chosen so that the low-order components
	# can never sum high enough to disturb a higher-order one, hence maximising the
	# packed sum == lexicographic optimisation of the three criteria in order.
	# Python's arbitrary-precision ints make this exact (no float drift).

	_SIM_SCALE = 10**9 # resolution of the primary score
	_SEC_SCALE = 10**6 # resolution of the secondary (shared-weight) score


	def _pack_scores(
	counts,
	sizes_a,
	sizes_b,
	metric,
	threshold,
	weight_total,
	priority_a=None,
	priority_b=None,
	combine=max,
	):
	"""Sparse dicts {(i, j): packed_int} and {(i, j): raw_similarity}.

	`threshold` is applied to the RAW similarity, before the priority
	multiplier -- otherwise a high-priority pair could sneak a structurally
	poor match past the cutoff.
	"""
	ka, kb = len(sizes_a), len(sizes_b)
	k_max = max(ka, kb, 1)

	tert_bound = k_max * k_max + k_max + 1 # tertiary < tert_bound
	base_sec = k_max * tert_bound # blocks tertiary carry
	sec_bound = int(weight_total * _SEC_SCALE) + 1 # secondary < sec_bound
	base_pri = base_sec * (k_max * sec_bound + 1) # blocks secondary carry

	packed: Dict[Tuple[int, int], int] = {}
	sims: Dict[Tuple[int, int], float] = {}
	for (i, j), c in counts.items():
	raw = metric(c, sizes_a[i], sizes_b[j])
	if raw < threshold or raw <= 0.0:
	continue
	effective = raw * _priority_of(priority_a, priority_b, i, j, combine)
	primary = int(round(effective * _SIM_SCALE))
	secondary = min(int(round(c * _SEC_SCALE)), sec_bound - 1) # shared weight
	tertiary = (k_max - i) * k_max + (k_max - j) # lower ids win
	packed[(i, j)] = primary * base_pri + secondary * base_sec + tertiary
	sims[(i, j)] = raw
	return packed, sims


	# --------------------------------------------------------------------------- #
	# 5. Assignment #
	# --------------------------------------------------------------------------- #


	def _hungarian(cost: List[List[int]]) -> List[int]:
	"""Min-cost assignment on a square integer matrix (Jonker-Volgenant style
	successive shortest paths). Returns row -> column.

	Strictly integral: the entries are arbitrary-precision packed scores that
	routinely exceed 2**53, so a float sentinel would risk silently collapsing
	distinct scores. `big` is derived from the matrix and provably exceeds any
	reduced cost -- potentials grow by at most one delta per iteration, so
	\|u\| + \|v\| <= 2k * max_abs, and cur = cost - u - v < (2k + 1) * max_abs.
	"""
	k = len(cost)
	if k == 0:
	return []
	max_abs = max((abs(x) for row in cost for x in row), default=0)
	big: int = (4 * k + 8) * (max_abs + 1)

	u: List[int] = [0] * (k + 1)
	v: List[int] = [0] * (k + 1)
	p: List[int] = [0] * (k + 1)
	way: List[int] = [0] * (k + 1)

	for i in range(1, k + 1):
	p[0] = i
	j0 = 0
	minv: List[int] = [big] * (k + 1)
	used: List[bool] = [False] * (k + 1)
	while True:
	used[j0] = True
	i0 = p[j0]
	delta: int = big
	j1 = 0
	for j in range(1, k + 1):
	if not used[j]:
	cur = cost[i0 - 1][j - 1] - u[i0] - v[j]
	if cur < minv[j]:
	minv[j], way[j] = cur, j0
	if minv[j] < delta:
	delta, j1 = minv[j], j
	for j in range(k + 1):
	if used[j]:
	u[p[j]] += delta
	v[j] -= delta
	else:
	minv[j] -= delta
	j0 = j1
	if p[j0] == 0:
	break
	while j0:
	j1 = way[j0]
	p[j0] = p[j1]
	j0 = j1

	row_to_col: List[int] = [-1] * k
	for j in range(1, k + 1):
	if p[j]:
	row_to_col[p[j] - 1] = j - 1
	return row_to_col


	def match_two(
	a: Partition,
	b: Partition,
	*,
	metric: str = "jaccard",
	threshold: float = 0.0,
	element_weights: Dict[Hashable, float] \| None = None,
	priority_a: Priorities \| None = None,
	priority_b: Priorities \| None = None,
	combine: str = "max",
	weight_total: float \| None = None,
	) -> Tuple[List[Tuple[int, int, float]], List[int], List[int]]:
	"""Optimal, deterministic 1-1 correspondence between two partitions.

	Returns (pairs, unmatched_a, unmatched_b). Pairs are (i, j, raw_similarity).
	The similarity reported is the raw, un-boosted one, so it stays comparable
	across pairs with different priorities.
	"""
	fn = METRICS[metric]
	comb = COMBINERS[combine]
	counts, sizes_a, sizes_b = contingency(a, b, element_weights)
	if weight_total is None:
	weight_total = max(sum(sizes_a), sum(sizes_b), 1.0)

	packed, sims = _pack_scores(
	counts,
	sizes_a,
	sizes_b,
	fn,
	threshold,
	weight_total,
	priority_a,
	priority_b,
	comb,
	)

	ka, kb = len(a), len(b)
	k = max(ka, kb)
	if k == 0:
	return [], [], []

	# square, padded. Padding and sub-threshold pairs score 0 -> stay unmatched
	cost = [[0] * k for _ in range(k)]
	for (i, j), score in packed.items():
	cost[i][j] = -score # maximise weight == minimise negative weight

	row_to_col = _hungarian(cost)

	pairs, matched_a, matched_b = [], set(), set()
	for i in range(ka):
	j = row_to_col[i]
	if j < kb and (i, j) in sims:
	pairs.append((i, j, sims[(i, j)]))
	matched_a.add(i)
	matched_b.add(j)
	pairs.sort(key=lambda t: (-t[2], t[0], t[1]))

	unmatched_a = [i for i in range(ka) if i not in matched_a]
	unmatched_b = [j for j in range(kb) if j not in matched_b]
	return pairs, unmatched_a, unmatched_b


	def greedy_match(
	a: Partition,
	b: Partition,
	*,
	metric: str = "jaccard",
	threshold: float = 0.0,
	element_weights: Dict[Hashable, float] \| None = None,
	priority_a: Priorities \| None = None,
	priority_b: Priorities \| None = None,
	combine: str = "max",
	):
	"""O(n log n) fallback for very large numbers of groups.
	Guaranteed within 1/2 of the optimum. Near-exact on similar partitions."""
	fn = METRICS[metric]
	comb = COMBINERS[combine]
	counts, sizes_a, sizes_b = contingency(a, b, element_weights)
	cand = []
	for (i, j), c in counts.items():
	raw = fn(c, sizes_a[i], sizes_b[j])
	if raw > threshold and raw > 0:
	eff = raw * _priority_of(priority_a, priority_b, i, j, comb)
	cand.append((-eff, -c, i, j, raw))
	cand.sort(key=lambda t: t[:4])
	used_a, used_b, pairs = set(), set(), []
	for _, _, i, j, raw in cand:
	if i not in used_a and j not in used_b:
	used_a.add(i)
	used_b.add(j)
	pairs.append((i, j, raw))
	unmatched_a = [i for i in range(len(sizes_a)) if i not in used_a]
	unmatched_b = [j for j in range(len(sizes_b)) if j not in used_b]
	return pairs, unmatched_a, unmatched_b


	# --------------------------------------------------------------------------- #
	# 6. Multi-partition alignment (reference-based) #
	# --------------------------------------------------------------------------- #


	@dataclass
	class Correspondence:
	"""One correspondence class: at most one group per partition."""

	members: Dict[int, int] = field(default_factory=dict) # partition -> group
	similarities: Dict[int, float] = field(default_factory=dict) # raw, vs reference

	def group_ids(self) -> List[GroupId]:
	return sorted(self.members.items())


	@dataclass
	class Alignment:
	reference: int
	classes: List[Correspondence]

	def render(self, partitions: Sequence[Partition]) -> str:
	lines = [f"reference partition = {self.reference}"]
	for n, cls in enumerate(self.classes):
	parts = []
	for g, idx in sorted(cls.members.items()):
	elems = sorted(map(str, partitions[g][idx]))
	sim = cls.similarities.get(g)
	tag = "" if sim is None else f"~{sim:.2f}"
	parts.append(f"G{g}[{','.join(elems)}]{tag}")
	lines.append(f" {n}: " + " <-> ".join(parts))
	return "\n".join(lines)


	def _pick_reference(partitions, metric, element_weights, priorities, combine) -> int:
	"""Medoid: the partition most similar to all the others (cheap greedy score)."""
	best, best_score = 0, -1.0
	for r in range(len(partitions)):
	total = 0.0
	for o in range(len(partitions)):
	if o == r:
	continue
	pairs, _, _ = greedy_match(
	partitions[r],
	partitions[o],
	metric=metric,
	element_weights=element_weights,
	priority_a=priorities[r] if priorities else None,
	priority_b=priorities[o] if priorities else None,
	combine=combine,
	)
	total += sum(s for _, _, s in pairs)
	if total > best_score:
	best, best_score = r, total
	return best


	def align(
	partitions: Sequence[Partition],
	*,
	reference: int \| str = "medoid",
	metric: str = "jaccard",
	threshold: float = 0.0,
	element_weights: Dict[Hashable, float] \| None = None,
	priorities: Sequence[Priorities] \| None = None,
	combine: str = "max",
	exact: bool = True,
	) -> Alignment:
	"""Align every partition to one reference partition.

	`element_weights` maps an element to its importance (default 1.0).
	`priorities` is parallel to `partitions`: one multiplier per group.

	Groups with no counterpart in the reference become their own correspondence
	class (like `[B]` in the example). Note the consequence of the
	reference-based strategy: two non-reference partitions that share an orphan
	group are not linked to each other, because all links go through the
	reference.
	"""
	if not partitions:
	return Alignment(0, [])
	if priorities is not None and len(priorities) != len(partitions):
	raise ValueError("`priorities` must have one entry per partition")

	if reference == "medoid":
	ref = _pick_reference(partitions, metric, element_weights, priorities, combine)
	else:
	ref = int(reference)

	# a single shared bound keeps the integer packing consistent across pairings
	weight_total = (
	max(
	(
	sum(
	element_weights.get(e, 1.0) if element_weights else 1.0
	for e in {x for grp in g for x in grp}
	)
	for g in partitions
	),
	default=1.0,
	)
	or 1.0
	)

	classes = [Correspondence(members={ref: i}) for i in range(len(partitions[ref]))]
	matcher = match_two if exact else greedy_match

	for g in range(len(partitions)):
	if g == ref:
	continue
	kwargs = dict(
	metric=metric,
	threshold=threshold,
	element_weights=element_weights,
	priority_a=priorities[ref] if priorities else None,
	priority_b=priorities[g] if priorities else None,
	combine=combine,
	)
	if exact:
	kwargs["weight_total"] = weight_total
	pairs, _, unmatched_other = matcher(partitions[ref], partitions[g], **kwargs)
	for i, j, sim in pairs:
	classes[i].members[g] = j
	classes[i].similarities[g] = sim
	for j in unmatched_other: # orphan -> own class
	classes.append(Correspondence(members={g: j}, similarities={g: 0.0}))

	return Alignment(ref, classes)


	# --------------------------------------------------------------------------- #
	# Demo #
	# --------------------------------------------------------------------------- #

	if __name__ == "__main__":
	partitions = [
	[["A", "B", "C", "D"], ["E", "F"]], # partition 0
	[["A", "C", "D"], ["E", "F"], ["B"]], # partition 1
	[["A", "C"], ["D", "B"], ["E", "F"]], # partition 2
	]

	print("--- unweighted (unchanged from before) ---")
	print(align(partitions[:2], reference=0).render(partitions[:2]))

	print("\n--- element weights: B is important, so [B] matters more ---")
	print(
	align(partitions[:2], reference=0, element_weights={"B": 10.0}).render(
	partitions[:2]
	)
	)

	# two groups in `right` compete for the same partner in `left`
	left = [["A", "B", "C"]]
	right = [["A", "B", "C", "D"], ["A"]]
	print("\n--- group priority overriding a better structural match ---")
	for pri in (None, [[1.0], [1.0, 3.0]]):
	pairs, _, _ = match_two(
	left,
	right,
	priority_a=pri[0] if pri else None,
	priority_b=pri[1] if pri else None,
	)
	i, j, s = pairs[0]
	print(
	f" priorities={pri}: -> right[{j}] "
	f"{sorted(right[j])} at raw similarity {s:.2f}"
	)
No results found