ОБНОВЛЕНИЕ:
Вот решение, которое генерирует только необходимые комбинации.Это в принципе быстрее, хотя сложность все еще экспоненциальна, и вы можете быстро выйти за пределы рекурсии.
def my_vectors(k, n):
# Minimum repetitions per element
base_repetitions = [n] * k
# "Unassigned" repetitions
rest = k ** n - k * n
# List reused for permutation construction
permutation = [-1] * (k ** n)
# For each possible repetition assignment
for repetitions in make_repetitions(base_repetitions, rest):
# Make all possible permutations
yield from make_permutations(repetitions, permutation)
# Finds all possible repetition assignments
def make_repetitions(repetitions, rest, first=0):
if rest <= 0:
yield repetitions
else:
for i in range(first, len(repetitions)):
repetitions[i] += 1
yield from make_repetitions(repetitions, rest - 1, i)
repetitions[i] -= 1
# Make all permutations with repetitions
def make_permutations(repetitions, permutation, idx=0):
if idx >= len(permutation):
yield list(permutation)
# If you are going to use the permutation within a loop only
# maybe you can avoid copying the list and do just:
# yield permutation
else:
for elem in range(len(repetitions)):
if repetitions[elem] > 0:
repetitions[elem] -= 1
permutation[idx] = elem
yield from make_permutations(repetitions, permutation, idx + 1)
repetitions[elem] += 1
for v in my_vectors(3, 2):
print(v)
Вывод:
(0, 0, 0, 0, 0, 1, 1, 2, 2)
(0, 0, 0, 0, 0, 1, 2, 1, 2)
(0, 0, 0, 0, 0, 1, 2, 2, 1)
(0, 0, 0, 0, 0, 2, 1, 1, 2)
(0, 0, 0, 0, 0, 2, 1, 2, 1)
(0, 0, 0, 0, 0, 2, 2, 1, 1)
(0, 0, 0, 0, 1, 0, 1, 2, 2)
(0, 0, 0, 0, 1, 0, 2, 1, 2)
(0, 0, 0, 0, 1, 0, 2, 2, 1)
(0, 0, 0, 0, 1, 1, 0, 2, 2)
...
Это неэффективно, нопростой способ реализовать это:
from itertools import product
from collections import Counter
def my_vectors(k, n):
for v in product(range(k), repeat=k ** n):
count = Counter(v)
if all(count[i] >= n for i in range(k)):
yield v
for v in my_vectors(3, 2):
print(v)
Вывод:
(0, 0, 0, 0, 0, 1, 1, 2, 2)
(0, 0, 0, 0, 0, 1, 2, 1, 2)
(0, 0, 0, 0, 0, 1, 2, 2, 1)
(0, 0, 0, 0, 0, 2, 1, 1, 2)
(0, 0, 0, 0, 0, 2, 1, 2, 1)
(0, 0, 0, 0, 0, 2, 2, 1, 1)
(0, 0, 0, 0, 1, 0, 1, 2, 2)
(0, 0, 0, 0, 1, 0, 2, 1, 2)
(0, 0, 0, 0, 1, 0, 2, 2, 1)
(0, 0, 0, 0, 1, 1, 0, 2, 2)
...
Очевидно, что, как только ваши числа станут немного больше, они будут работать вечно, так что это полезно только дляочень маленькие проблемы или в качестве базового уровня для сравнения.
В любом случае, количество элементов, которые создает проблема, в любом случае экспоненциально велико, поэтому, хотя вы можете сделать это значительно лучше (т.е. генерировать только правильные элементы вместовсе возможные и отбрасываемые), он не может быть "быстрым" для любого размера.