构造所有带有两个非零的`r`元组

Question

给定一个 int 值

val

和一个元组长度

r

，我需要创建所有

r

元组，其中

d

为

{+val, -val}

，其余的用零填充。有了

d=2

，我就能做到

val = 7
r = 5

out = []
for i0 in range(r - 1):
    for i1 in range(i0 + 1, r):
        for v0 in [+val, -val]:
            for v1 in [+val, -val]:
                t = [0] * r
                t[i0] = v0
                t[i1] = v1
                print(t)

[7, 7, 0, 0, 0]
[7, -7, 0, 0, 0]
[-7, 7, 0, 0, 0]
[-7, -7, 0, 0, 0]
[7, 0, 7, 0, 0]
# ...

但这已经让人感觉很混乱了。对于更大的

d

来说，情况变得更糟。我查看了 itertools 组合迭代器，但这些似乎都没有帮助。

有什么提示吗？

Answer 1

使用：

• combinations_with_replacement((-val, +val), d) 获取 val 和 -val 的所有无序 d-uplet；
• more_itertools.distinct_permutations 以所有可能的方式用（r-d）零填充。

from itertools import combinations_with_replacement, chain
from more_itertools import distinct_permutations

def combs_with_zeroes(val, d, r):
    assert r >= d
    return chain.from_iterable(distinct_permutations((*tuple_d, *(0,)*(r-d))) for tuple_d in combinations_with_replacement((-val, +val), d))

print(list(combs_with_zeroes(val=7, d=1, r=2)))
# [(-7, 0), (0, -7), (0, 7), (7, 0)]

print(list(combs_with_zeroes(val=7, d=1, r=3)))
# [(-7, 0, 0), (0, -7, 0), (0, 0, -7), (0, 0, 7), (0, 7, 0), (7, 0, 0)]

print(list(combs_with_zeroes(val=7, d=3, r=4)))
# [(-7, -7, -7, 0), (-7, -7, 0, -7), (-7, 0, -7, -7), (0, -7, -7, -7), (-7, -7, 0, 7), (-7, -7, 7, 0), (-7, 0, -7, 7), (-7, 0, 7, -7), (-7, 7, -7, 0), (-7, 7, 0, -7), (0, -7, -7, 7), (0, -7, 7, -7), (0, 7, -7, -7), (7, -7, -7, 0), (7, -7, 0, -7), (7, 0, -7, -7), (-7, 0, 7, 7), (-7, 7, 0, 7), (-7, 7, 7, 0), (0, -7, 7, 7), (0, 7, -7, 7), (0, 7, 7, -7), (7, -7, 0, 7), (7, -7, 7, 0), (7, 0, -7, 7), (7, 0, 7, -7), (7, 7, -7, 0), (7, 7, 0, -7), (0, 7, 7, 7), (7, 0, 7, 7), (7, 7, 0, 7), (7, 7, 7, 0)]

print(list(combs_with_zeroes(val=7, d=2, r=5)))
# [(-7, -7, 0, 0, 0), (-7, 0, -7, 0, 0), (-7, 0, 0, -7, 0), (-7, 0, 0, 0, -7), (0, -7, -7, 0, 0), (0, -7, 0, -7, 0), (0, -7, 0, 0, -7), (0, 0, -7, -7, 0), (0, 0, -7, 0, -7), (0, 0, 0, -7, -7), (-7, 0, 0, 0, 7), (-7, 0, 0, 7, 0), (-7, 0, 7, 0, 0), (-7, 7, 0, 0, 0), (0, -7, 0, 0, 7), (0, -7, 0, 7, 0), (0, -7, 7, 0, 0), (0, 0, -7, 0, 7), (0, 0, -7, 7, 0), (0, 0, 0, -7, 7), (0, 0, 0, 7, -7), (0, 0, 7, -7, 0), (0, 0, 7, 0, -7), (0, 7, -7, 0, 0), (0, 7, 0, -7, 0), (0, 7, 0, 0, -7), (7, -7, 0, 0, 0), (7, 0, -7, 0, 0), (7, 0, 0, -7, 0), (7, 0, 0, 0, -7), (0, 0, 0, 7, 7), (0, 0, 7, 0, 7), (0, 0, 7, 7, 0), (0, 7, 0, 0, 7), (0, 7, 0, 7, 0), (0, 7, 7, 0, 0), (7, 0, 0, 0, 7), (7, 0, 0, 7, 0), (7, 0, 7, 0, 0), (7, 7, 0, 0, 0)]

注意结果变得相当长而且很快；元素总数为

2**d * (r choose d)

。