ANF algorithm in Python
Estimates neighbourhood functions with Flajolet-Martin style sketches.
- File
anf_algorithm.py, 58 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/anf-algorithm/. The authors are listed in prov/MDNVAV.jsonl.
import random
PHI = 0.77351 # the Flajolet-Martin correction constant
def random_ranks(n: int, r: int, seed: int) -> list[list[int]]:
"""For each vertex and each of r copies, a bit index i with probability 2^-(i+1)."""
rng = random.Random(seed)
ranks: list[list[int]] = []
for _ in range(n):
row = []
for _ in range(r):
i = 0
while rng.random() < 0.5 and i < 62:
i += 1
row.append(i)
ranks.append(row)
return ranks
def lowest_zero(mask: int) -> int:
"""The index of the lowest bit of mask that is 0."""
return ((mask + 1) & ~mask).bit_length() - 1
def fm_estimate(masks: list[int]) -> float:
"""The Flajolet-Martin estimate of a set size from r bitmasks of the set."""
mean = sum(lowest_zero(x) for x in masks) / len(masks)
return 2.0 ** mean / PHI
def anf_algorithm(adj: list[list[int]], h_max: int, ranks: list[list[int]]) -> tuple[list[float], list[list[int]]]:
"""Estimate the neighbourhood function N(0), ..., N(h_max) of a graph.
adj[v] lists the out-neighbours of v (both directions for an undirected graph);
N(h) is the number of ordered pairs (v, u) with u reachable from v by at most h
arcs, v = u included. ranks[v][j] is the random bit of v in copy j (see
random_ranks). Returns the estimates and the final bitmasks of each vertex.
Time O(h_max * r * (n + m)), space O(n * r) machine words.
"""
n = len(adj)
r = len(ranks[0]) if n else 0
masks = [[1 << ranks[v][j] for j in range(r)] for v in range(n)] # @line L1
est = [sum(fm_estimate(masks[v]) for v in range(n))] # @line L2
for _ in range(h_max): # @line L3
new = [list(masks[v]) for v in range(n)] # @line L4
for v in range(n): # @line L5
row = new[v]
for u in adj[v]:
other = masks[u]
for j in range(r):
row[j] |= other[j]
masks = new # @line L6
est.append(sum(fm_estimate(masks[v]) for v in range(n))) # @line L7
return est, masks # @line L8
test_anf_algorithm.py 105 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/anf-algorithm/. The authors are listed in prov/MDNVAV.jsonl.
import random
from collections import deque
from anf_algorithm import anf_algorithm, fm_estimate, lowest_zero, random_ranks
def undirected(n: int, edges: list[tuple[int, int]]) -> list[list[int]]:
adj: list[list[int]] = [[] for _ in range(n)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
return adj
def ball(adj: list[list[int]], v: int, h: int) -> list[int]:
"""The vertices reachable from v by at most h arcs, by breadth-first search."""
dist = {v: 0}
q = deque([v])
while q:
x = q.popleft()
if dist[x] == h:
continue
for y in adj[x]:
if y not in dist:
dist[y] = dist[x] + 1
q.append(y)
return list(dist)
def direct_masks(adj: list[list[int]], h: int, ranks: list[list[int]]) -> list[list[int]]:
n = len(adj)
r = len(ranks[0]) if n else 0
out = []
for v in range(n):
row = [0] * r
for u in ball(adj, v, h):
for j in range(r):
row[j] |= 1 << ranks[u][j]
out.append(row)
return out
EXAMPLE = [(0, 1), (0, 2), (1, 2), (1, 3), (2, 3), (3, 4), (4, 5), (4, 6), (5, 6)]
def test_helpers() -> None:
assert lowest_zero(0) == 0
assert lowest_zero(0b1011) == 2
assert lowest_zero(0b111) == 3
assert abs(fm_estimate([0b1, 0b1]) - 2 / 0.77351) < 1e-12
def test_example() -> None:
adj = undirected(7, EXAMPLE)
exact = [sum(len(ball(adj, v, h)) for v in range(7)) for h in range(5)]
assert exact == [7, 25, 35, 45, 49]
ranks = random_ranks(7, 64, 1)
est, masks = anf_algorithm(adj, 4, ranks)
assert masks == direct_masks(adj, 4, ranks)
assert len(est) == 5
assert all(a <= b for a, b in zip(est, est[1:]))
def test_edges() -> None:
assert anf_algorithm([], 3, []) == ([0, 0, 0, 0], [])
est, masks = anf_algorithm([[]], 2, [[0, 1]])
assert masks == [[1, 2]]
assert est[0] == est[1] == est[2]
# a directed path 0 -> 1 -> 2: vertex 0 reaches everything, vertex 2 only itself
ranks = [[0], [1], [2]]
_, masks = anf_algorithm([[1], [2], []], 2, ranks)
assert masks == [[0b111], [0b110], [0b100]]
def test_accuracy_on_a_large_graph() -> None:
rng = random.Random(77)
n = 1500
edges = [(i, i + 1) for i in range(n - 1)] + [(rng.randrange(n), rng.randrange(n)) for _ in range(300)]
adj = undirected(n, edges)
est, _ = anf_algorithm(adj, 6, random_ranks(n, 64, 5))
for h in (3, 6):
exact = sum(len(ball(adj, v, h)) for v in range(n))
assert abs(est[h] - exact) < 0.1 * exact
def test_random_against_oracle() -> None:
rng = random.Random(4242)
for _ in range(500):
n = rng.randint(1, 9)
directed = rng.random() < 0.5
adj: list[list[int]] = [[] for _ in range(n)]
for _ in range(rng.randint(0, 2 * n)):
u, v = rng.randrange(n), rng.randrange(n)
adj[u].append(v)
if not directed:
adj[v].append(u)
h = rng.randint(0, 4)
ranks = random_ranks(n, rng.randint(1, 8), rng.randrange(10**6))
est, masks = anf_algorithm(adj, h, ranks)
want = direct_masks(adj, h, ranks)
assert masks == want
assert abs(est[-1] - sum(fm_estimate(row) for row in want)) < 1e-9
What the next level needs
- the tests did not run in this build