用 0 和 1 填充的 numpy 二维数组的所有组合

Posted 2023-03-12

【中文标题】用 0 和 1 填充的 numpy 二维数组的所有组合

【英文标题】：All combinations of a numpy 2d array filled with 0s and 1s

【发布时间】：2022-01-09 16:08:16

【问题描述】：

给定 K，我需要拥有 K x 2 numpy 矩阵的所有可能组合，以便在每个矩阵中，除了不同行和列中的两个 1 之外，所有的都是 0。 对于 K = 5 是这样的：

[[1,0],[0,1],[0,0],[0,0][0,0]] [[1,0],[0,0],[0,1],[0,0][0,0]] [[1,0],[0,0],[0,0],[0,1][0,0]] [[1,0],[0,0],[0,0],[0,0][0,1]] [[0,0],[1,0],[0,1],[0,0][0,0]] [[0,0],[1,0],[0,0],[0,1][0,0]] ...等等

所以结果数组应该是 K x 2 x (K*(K-1)/2)。 我想避免循环，因为当 K 足够大时（在我的特定情况下 K = 300），这不是一种有效的方法

【参考方案1】：

我想不出一个优雅的解决方案，但这里有一个不太优雅的纯 numpy 解决方案：

import numpy as np

def combination_matrices(K):
    # get combination indices
    i, j = np.indices((K, K))
    comb_indices = np.transpose((i < j).nonzero())  # (num_combs, 2) array where ones are
    num_combs = comb_indices.shape[0]  # K*(K-1)/2

    # create a matrix of the desired shape, first axis enumerates combinations
    matrices = np.zeros((num_combs, K, 2), dtype=int)
    # broadcasting assignment of ones
    comb_range, col_index = np.ogrid[:num_combs, :2]
    matrices[comb_range, comb_indices, col_index] = 1
    return matrices

这首先使用(K, K) 形数组的索引来查找每个组合的索引对（这些索引对数组的上三角形进行编码，不包括对角线）。然后我们使用有点棘手的广播赋值（重 fancy indexing）将预分配的输出数组的每个对应元素设置为 1。

请注意，我将 K*(K-1)/2 大小的轴放在首位，因为这在具有 C 连续内存布局的 numpy 中最有意义。这样，当您将矩阵用于组合索引 3 时，arr[3, ...] 将是形状为 (K, 2) 的连续内存块，在矢量化操作中可以快速使用。

K = 4 的输出：

[[[1 0]
  [0 1]
  [0 0]
  [0 0]]

 [[1 0]
  [0 0]
  [0 1]
  [0 0]]

 [[1 0]
  [0 0]
  [0 0]
  [0 1]]

 [[0 0]
  [1 0]
  [0 1]
  [0 0]]

 [[0 0]
  [1 0]
  [0 0]
  [0 1]]

 [[0 0]
  [0 0]
  [1 0]
  [0 1]]]

【讨论】：

谢谢！这正是我需要的。我只有一个问题：我需要沿 K 轴访问这个数组，所以将 K 作为第一个维度可能会更好？我不是内存分配方面的专家

@mariottidae 在这种情况下，您可能只需将分配更改为matrices = np.zeros((K, 2, num_combs), dtype=int) 并将分配更改为matrices[comb_indices, col_index, comb_range] = 1。这将为您提供一个形状数组(K, 2, num_combs)。相当于在我的原始版本上调用.transpose(1, 2, 0)，但这将再次是内存中的一个连续数组（与转置不同）。

【参考方案2】：

这是一个奇怪的具体问题，但是一个有趣的问题，我很想知道上下文是什么？

您正在寻找多重集合 的所有排列，python 的 itertools 目前不支持此功能。所以最简单的解决方案是使用sympy 库的多集工具。

以下代码在我的机器上运行大约需要 2.5 分钟，因此对于单线程来说相当快。您正在查看 K=300 的 179700 唯一排列。

（我的灵感来自https://***.com/a/40289807/10739860）

from collections import Counter
from math import factorial, prod

import numpy as np
from sympy.utilities.iterables import multiset_permutations
from tqdm import tqdm


def No_multiset_permutations(multiset: list) -> int:
    """Calculates the No. possible permutations given a multiset.
    See: https://en.wikipedia.org/wiki/Permutation#Permutations_of_multisets

    :param multiset: List representing a multiset.
    """
    value_counts = Counter(multiset).values()
    denominator = prod([factorial(val) for val in value_counts])
    return int(factorial(len(multiset)) / denominator)


def multiset_Kx2_permutations(K: int) -> np.ndarray:
    """This will generate all possible unique Kx2 permutations of an array
    withsize K where two values are 1 and the rest are 0.

    :param K: The size of the array.
    """
    # Construct number multiset, e.g. K=5 gives [1, 1, 0, 0, 0, 0, 0, 0, 0, 0]
    numbers = [1, 1] + [0] * (K - 1) * 2

    # Use sympy's multiset_permutations to get a multiset permutation generator
    generator = multiset_permutations(numbers)

    # Calculate the No. possible permutations
    number_of_perms = No_multiset_permutations(numbers)

    # Get all permutations, bonus progress bar is included :)
    unique_perms = [next(generator) for _ in tqdm(range(number_of_perms))]

    # Reshape each permutation to Kx2
    unique_perms = np.array(unique_perms, dtype=np.int8)
    return unique_perms.reshape(-1, K, 2)


if __name__ == "__main__":
    solution = multiset_Kx2_permutations(300)

【参考方案3】：

另一种可能性（重新排列轴以获得更清晰的输出）：

from itertools import combinations
import numpy as np

k = 4
x = list(combinations(range(k), 2))
out = np.zeros((n := len(x), k, 2), dtype=int)
out[np.c_[:n], x, [0, 1]] = 1
print(out)

它给出：

[[[1 0]
  [0 1]
  [0 0]
  [0 0]]

 [[1 0]
  [0 0]
  [0 1]
  [0 0]]

 [[1 0]
  [0 0]
  [0 0]
  [0 1]]

 [[0 0]
  [1 0]
  [0 1]
  [0 0]]

 [[0 0]
  [1 0]
  [0 0]
  [0 1]]

 [[0 0]
  [0 0]
  [1 0]
  [0 1]]]