Python/backtracking/all_combinations.py
pre-commit-ci[bot] bc8df6de31
[pre-commit.ci] pre-commit autoupdate (#11322)
* [pre-commit.ci] pre-commit autoupdate

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.2.2 → v0.3.2](https://github.com/astral-sh/ruff-pre-commit/compare/v0.2.2...v0.3.2)
- [github.com/pre-commit/mirrors-mypy: v1.8.0 → v1.9.0](https://github.com/pre-commit/mirrors-mypy/compare/v1.8.0...v1.9.0)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2024-03-13 07:52:41 +01:00

83 lines
2.4 KiB
Python

"""
In this problem, we want to determine all possible combinations of k
numbers out of 1 ... n. We use backtracking to solve this problem.
Time complexity: O(C(n,k)) which is O(n choose k) = O((n!/(k! * (n - k)!))),
"""
from __future__ import annotations
from itertools import combinations
def combination_lists(n: int, k: int) -> list[list[int]]:
"""
>>> combination_lists(n=4, k=2)
[[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
"""
return [list(x) for x in combinations(range(1, n + 1), k)]
def generate_all_combinations(n: int, k: int) -> list[list[int]]:
"""
>>> generate_all_combinations(n=4, k=2)
[[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
>>> generate_all_combinations(n=0, k=0)
[[]]
>>> generate_all_combinations(n=10, k=-1)
Traceback (most recent call last):
...
ValueError: k must not be negative
>>> generate_all_combinations(n=-1, k=10)
Traceback (most recent call last):
...
ValueError: n must not be negative
>>> generate_all_combinations(n=5, k=4)
[[1, 2, 3, 4], [1, 2, 3, 5], [1, 2, 4, 5], [1, 3, 4, 5], [2, 3, 4, 5]]
>>> from itertools import combinations
>>> all(generate_all_combinations(n, k) == combination_lists(n, k)
... for n in range(1, 6) for k in range(1, 6))
True
"""
if k < 0:
raise ValueError("k must not be negative")
if n < 0:
raise ValueError("n must not be negative")
result: list[list[int]] = []
create_all_state(1, n, k, [], result)
return result
def create_all_state(
increment: int,
total_number: int,
level: int,
current_list: list[int],
total_list: list[list[int]],
) -> None:
if level == 0:
total_list.append(current_list[:])
return
for i in range(increment, total_number - level + 2):
current_list.append(i)
create_all_state(i + 1, total_number, level - 1, current_list, total_list)
current_list.pop()
if __name__ == "__main__":
from doctest import testmod
testmod()
print(generate_all_combinations(n=4, k=2))
tests = ((n, k) for n in range(1, 5) for k in range(1, 5))
for n, k in tests:
print(n, k, generate_all_combinations(n, k) == combination_lists(n, k))
print("Benchmark:")
from timeit import timeit
for func in ("combination_lists", "generate_all_combinations"):
print(f"{func:>25}(): {timeit(f'{func}(n=4, k = 2)', globals=globals())}")