AllPairs in Python
Finds all similar pairs by indexing prefix tokens in an inverted index and verifying candidates.
- File
allpairs.py, 64 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/allpairs/. The authors are listed in prov/0NZTCD.jsonl.
from collections import Counter, defaultdict
from collections.abc import Sequence
SparseVector = dict[int, float]
def dot(x: SparseVector, y: SparseVector) -> float:
if len(x) > len(y):
x, y = y, x
return sum(w * y[i] for i, w in x.items() if i in y)
def all_pairs(vectors: Sequence[SparseVector], t: float) -> list[tuple[int, int, float]]:
"""Return all (j, i, x_j . x_i) with j < i and dot product >= t, for nonnegative sparse vectors.
All-Pairs of Bayardo, Ma and Srikant with the max-weight bound over the whole input.
Time: O(N log N + C + sum of verified prefix lengths), where C is the number of index entries scanned.
"""
if t <= 0:
raise ValueError("the threshold must be positive")
freq = Counter(i for x in vectors for i, w in x.items() if w > 0)
rank = {i: r for r, (i, _) in enumerate(sorted(freq.items(), key=lambda kv: (-kv[1], kv[0])))}
maxw: dict[int, float] = defaultdict(float) # @line L1
for x in vectors:
for i, w in x.items():
maxw[i] = max(maxw[i], w)
out: list[tuple[int, int, float]] = [] # @line L2
index: dict[int, list[tuple[int, float]]] = defaultdict(list)
prefix: list[SparseVector] = []
for xid, x in enumerate(vectors): # @line L3
out.extend((y, xid, s) for y, s in find_matches(x, index, prefix, t)) # @line L4
b = 0.0 # @line L5
unindexed: SparseVector = {}
for i in sorted((i for i, w in x.items() if w > 0), key=rank.__getitem__): # @line L6
b += maxw[i] * x[i] # @line L7
if b >= t: # @line L8
index[i].append((xid, x[i]))
else:
unindexed[i] = x[i]
prefix.append(unindexed)
return out # @line L9
def find_matches(
x: SparseVector,
index: dict[int, list[tuple[int, float]]],
prefix: list[SparseVector],
t: float,
) -> list[tuple[int, float]]:
"""Return (y, x . y) for each earlier vector y with x . y >= t."""
acc: dict[int, float] = defaultdict(float) # @line L10
for i, w in x.items(): # @line L11
for y, wy in index.get(i, ()):
acc[y] += w * wy
found = []
for y, a in acc.items(): # @line L12
if a > 0:
s = a + dot(x, prefix[y])
if s >= t:
found.append((y, s))
return found # @line L13
test_allpairs.py 58 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/allpairs/. The authors are listed in prov/0NZTCD.jsonl.
import random
from allpairs import all_pairs, dot
EXAMPLE = [
{0: 0.6, 1: 0.8},
{0: 0.8, 2: 0.6},
{0: 0.36, 1: 0.48, 3: 0.8},
{2: 0.6, 3: 0.8},
]
def brute(vs, t):
return sorted((j, i) for i in range(len(vs)) for j in range(i) if dot(vs[i], vs[j]) >= t)
def test_example() -> None:
got = all_pairs(EXAMPLE, 0.55)
assert [(a, b) for a, b, _ in got] == [(0, 2), (2, 3)]
assert abs(got[0][2] - 0.6) < 1e-12 and abs(got[1][2] - 0.64) < 1e-12
def test_rejects_zero_threshold() -> None:
try:
all_pairs(EXAMPLE, 0)
except ValueError:
return
raise AssertionError("expected ValueError")
def test_random_integer_vectors() -> None:
rng = random.Random(4242)
for _ in range(600):
n, d = rng.randint(0, 25), rng.randint(1, 12)
vs = []
for _ in range(n):
vs.append({i: rng.randint(1, 5) for i in range(d) if rng.random() < rng.choice((0.15, 0.4, 0.8))})
t = rng.randint(1, 40)
got = all_pairs(vs, t)
assert sorted((a, b) for a, b, _ in got) == brute(vs, t)
assert all(s == dot(vs[a], vs[b]) for a, b, s in got)
def test_random_sets_cosine() -> None:
rng = random.Random(77)
for _ in range(100):
sets = [frozenset(rng.sample(range(15), rng.randint(1, 6))) for _ in range(30)]
vs = [{i: 1 / len(s) ** 0.5 for i in s} for s in sets]
t = rng.choice((0.3, 0.5, 0.7, 0.9))
got = {(a, b) for a, b, _ in all_pairs(vs, t)}
for i in range(30):
for j in range(i):
c = len(sets[i] & sets[j]) / (len(sets[i]) * len(sets[j])) ** 0.5
if abs(c - t) > 1e-9:
assert ((j, i) in got) == (c >= t)
What the next level needs
- the tests did not run in this build