如何在Python中快速生成这些向量?

时间:2016-04-22 00:21:24

标签: python

假设我想要生成长度为n的所有向量,使得k位置的n个成分正好是+1或-1(对它的限制没有限制+ 1或-1),其余的nk分量均为零。

我可以按如下方式编写代码:

from itertools import combinations, product
result = []    
for x in combinations(range(n),k):
    for y in product([-1,1],repeat=k)
        zero = [0] * n
        for a in x:
            for b in y:
                zero[a] = b
                result.append(zero)

这种方式有效,但我认为这有点乏味。有没有快速的方法来给出结果?

2 个答案:

答案 0 :(得分:2)

让我们编写一个函数,生成长test test n 1s和k 0s的所有向量:

n - k

然后,让我们把一些1转成-1s:

def gen(n, k):
    for indices in itertools.combinations(range(n), k):
        l = [0] * n
        for i in indices:
            l[i] = 1
        yield(l)

示例输出:

def gen(n, k):
    for indices in itertools.combinations(range(n), k):
        # each one of indices can be -1 or +1
        for i in range(len(indices) + 1):
            for neg_indices in itertools.combinations(indices, i):
                l = [0] * n
                for i in indices:
                    l[i] = 1
                for i in neg_indices:
                    l[i] = -1
                yield(l)

您的原始实施相当接近;解决它:

>>> list(gen(3, 2))
[[1, 1, 0], [-1, 1, 0], [1, -1, 0], [-1, -1, 0], [1, 0, 1], [-1, 0, 1], [1, 0, -1], [-1, 0, -1], [0, 1, 1], [0, -1, 1], [0, 1, -1], [0, -1, -1]]

请注意使用def gen(n, k): for x in itertools.combinations(range(n), k): for y in itertools.product([-1,1], repeat=k): zero = [0] * n for a, b in zip(x, y): zero[a] = b yield zero 而不是嵌套。

答案 1 :(得分:0)

你需要的只是

import random
import itertools

# create a single sample vector containing the distribution of elements
base_list = map(lambda x: random.choice([1,-1]), range(k)) + [0]*(n-k)

# generate all unique permutations of this list
result = set(itertools.permutations(base_list))

result现在将包含base_list中包含的基本向量值的所有唯一排列。

结果排列都保证k元素为-1或1,n-k元素全为零。

set()负责确保重复的条目 - 因为permutations()按索引而不是值进行排列 - 被排除在外。

对于大型n,这种实现无疑会变慢,但对于小型列表应该足够了。