Two pointers with distinct counts in Python
Advance the right pointer, shrink the left pointer while the distinct count exceeds k, and track the best length.
- File
distinct_count_two_pointers.py, 71 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/distinct-count-two-pointers/. The authors are listed in prov/9MW95T.jsonl.
from collections.abc import Hashable, Sequence
def longest_at_most_k_distinct(a: Sequence[Hashable], k: int) -> tuple[int, int]:
"""Return (length, start) of a longest subarray a[start:start+length] with at most k distinct values.
Precondition: the items are hashable and k >= 0 (raises ValueError if k < 0).
For the empty answer the start is 0.
Cost: O(n) expected time over the hash table, O(min(n, k)) extra space.
"""
if k < 0:
raise ValueError("k must be non-negative")
cnt: dict[Hashable, int] = {} # @line L1
d = 0
i = 0
best = 0
start = 0
for j in range(len(a)): # @line L2
x = a[j]
if x not in cnt: # @line L3
cnt[x] = 0
d += 1
cnt[x] += 1 # @line L4
while d > k: # @line L5
y = a[i] # @line L6
cnt[y] -= 1
if cnt[y] == 0:
del cnt[y]
d -= 1
i += 1 # @line L7
if j + 1 - i > best: # @line L8
best = j + 1 - i
start = i
return best, start # @line L9
def count_at_most_k_distinct(a: Sequence[Hashable], k: int) -> int:
"""Return the number of non-empty subarrays with at most k distinct values.
Precondition: hashable items. Returns 0 for k <= 0.
Cost: O(n) expected time, O(min(n, k)) extra space.
"""
if k <= 0:
return 0
cnt: dict[Hashable, int] = {}
i = 0
total = 0
for j, x in enumerate(a):
cnt[x] = cnt.get(x, 0) + 1
while len(cnt) > k:
y = a[i]
cnt[y] -= 1
if cnt[y] == 0:
del cnt[y]
i += 1
total += j + 1 - i
return total
def count_exactly_k_distinct(a: Sequence[Hashable], k: int) -> int:
"""Return the number of non-empty subarrays with exactly k distinct values.
Precondition: hashable items. Returns 0 for k <= 0.
Cost: two passes, O(n) expected time.
"""
if k <= 0:
return 0
return count_at_most_k_distinct(a, k) - count_at_most_k_distinct(a, k - 1)
test_distinct_count_two_pointers.py 69 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/distinct-count-two-pointers/. The authors are listed in prov/9MW95T.jsonl.
import random
import pytest
from distinct_count_two_pointers import (
count_at_most_k_distinct,
count_exactly_k_distinct,
longest_at_most_k_distinct,
)
def brute_longest(a: list[int], k: int) -> int:
n = len(a)
return max(
(j - i for i in range(n + 1) for j in range(i, n + 1) if len(set(a[i:j])) <= k),
default=0,
)
def brute_count(a: list[int], k: int, exact: bool) -> int:
n = len(a)
total = 0
for i in range(n):
for j in range(i + 1, n + 1):
d = len(set(a[i:j]))
if (d == k) if exact else (d <= k):
total += 1
return total
def test_example() -> None:
assert longest_at_most_k_distinct([1, 2, 1, 2, 3, 3, 4, 1], 2) == (4, 0)
assert longest_at_most_k_distinct([1, 2, 3, 2, 2], 2) == (4, 1)
big = [10**9, -5, 10**9, 3, -5, -5, 3]
assert longest_at_most_k_distinct(big, 2) == (4, 3)
def test_counts() -> None:
assert count_at_most_k_distinct([1, 2, 1, 2, 3], 2) == 12
assert count_at_most_k_distinct([1, 2, 1, 2, 3], 1) == 5
assert count_exactly_k_distinct([1, 2, 1, 2, 3], 2) == 7
assert count_exactly_k_distinct([1, 2, 1, 2, 3], 0) == 0
def test_edges() -> None:
assert longest_at_most_k_distinct([], 3) == (0, 0)
assert longest_at_most_k_distinct([5, 5, 5], 0) == (0, 0)
assert longest_at_most_k_distinct([5, 5, 5], 1) == (3, 0)
assert longest_at_most_k_distinct([1, 2, 3], 5) == (3, 0)
assert longest_at_most_k_distinct("aabbbc", 1) == (3, 2)
with pytest.raises(ValueError):
longest_at_most_k_distinct([1], -1)
def test_against_brute_force() -> None:
rng = random.Random(20261105)
for _ in range(2000):
n = rng.randrange(0, 13)
a = [rng.randrange(0, rng.randrange(1, 6)) for _ in range(n)]
k = rng.randrange(0, 5)
length, start = longest_at_most_k_distinct(a, k)
assert length == brute_longest(a, k)
assert len(set(a[start:start + length])) <= k
assert 0 <= start <= n - length
assert count_at_most_k_distinct(a, k) == brute_count(a, k, False)
assert count_exactly_k_distinct(a, k) == brute_count(a, k, True)
What the next level needs
- the tests did not run in this build