beam: generalized Boyer-Moore majority vote with verification

wizard-work/beam/majority.py

@@ -0,0 +1,107 @@
+"""Generalized Boyer-Moore majority vote.
+
+Finds all elements occurring more than n/k times in a sequence, in O(n*k)
+time and O(k) extra space. Candidate selection alone is not sufficient:
+a final verification pass is mandatory, because candidates can be false
+positives produced by cancellation bookkeeping.
+"""
+from collections import Counter
+from typing import Iterable, List
+
+
+def majority_element_v2(nums):
+    data = list(nums)
+    cand = None
+    count = 0
+    for x in data:
+        if count == 0:
+            cand = x
+            count = 1
+        elif x == cand:
+            count += 1
+        else:
+            count -= 1
+    if not data:
+        return None
+    if data.count(cand) > len(data) // 2:
+        return cand
+    return None
+
+
+def majority_element(nums):
+    """Return the element occurring more than len(nums)//2 times, or None."""
+    return majority_element_v2(nums)
+
+
+def majority_gt_k(nums, k: int) -> List:
+    """Return all elements occurring more than len(nums)//k times (k>=2)."""
+    if k < 2:
+        raise ValueError("k must be >= 2")
+    data = list(nums)
+    n = len(data)
+    if n == 0:
+        return []
+    cand = {}
+    for x in data:
+        if x in cand:
+            cand[x] += 1
+        elif len(cand) < k - 1:
+            cand[x] = 1
+        else:
+            to_remove = []
+            for key in list(cand):
+                cand[key] -= 1
+                if cand[key] == 0:
+                    to_remove.append(key)
+            for key in to_remove:
+                del cand[key]
+    candidates = list(cand.keys())
+    threshold = n // k
+    counts = Counter(data)
+    result = []
+    seen = set()
+    for c in candidates:
+        if counts.get(c, 0) > threshold and c not in seen:
+            seen.add(c)
+            result.append(c)
+    return result
+
+
+def majority(nums):
+    return majority_element_v2(nums)
+
+
+def majority_third(nums):
+    return majority_gt_k(nums, 3)
+
+
+if __name__ == "__main__":
+    import random
+    rng = random.Random(42)
+    for trial in range(2000):
+        n = rng.randint(0, 40)
+        if n == 0:
+            data = []
+        else:
+            if rng.random() < 0.5:
+                maj = rng.randint(0, 5)
+                need = n // 2 + 1
+                data = [maj] * need + [rng.randint(0, 5) for _ in range(n - need)]
+                rng.shuffle(data)
+            else:
+                data = [rng.randint(0, 5) for _ in range(n)]
+        c = Counter(data)
+        bf = [x for x, v in c.items() if v > n // 2] if data else []
+        bf_majority = bf[0] if bf else None
+        got = majority_element_v2(data)
+        assert got == bf_majority, (data, got, bf_majority)
+        bf3 = [x for x, v in c.items() if v > n // 3]
+        got3 = majority_gt_k(data, 3)
+        assert sorted(got3) == sorted(bf3), (data, got3, bf3)
+    print("ALL RANDOMIZED CHECKS PASSED")
+    assert majority_element_v2([]) is None
+    assert majority_element_v2([1]) == 1
+    assert majority_element_v2([1, 2]) is None
+    assert majority_gt_k([], 3) == []
+    assert majority_gt_k([1, 1, 2, 2, 3, 3], 3) == []
+    print("EDGE CASES PASSED")