python - 通过meshgrid查找成对的numpy数组的所有组合

假设我有一个由值对组成的numpy数组。我想找到这些组合的所有组合，而又不会撕裂它们。特别是，我希望为此提供一个numpy.meshgrid解决方案。

想象一个构造如下的数组：

ab = np.array([[1,10], [2,20], [3,30], [4,40]])

然后我想要的输出是

>>> out: ([1,10], [2,20])
         ([1,10], [3,30])
         ([1,10], [4,40])
         ([2,20], [3,30])
         ([2,20], [4,40])
         ([3,30], [4,40])

输出可以是np.array或tuple（之后我可以进行相应转换）。请注意，结果中的重复项是如何省略的，而忽略了我夫妇的顺序（如果[[1,10], [2,20]]已经存在，则我不想在输出中使用[[2,20], [1,10]]）。对于实际情况，ab的大小为30,000，因此速度是另一个问题。

这就是为什么我首先尝试使用meshgrid的原因。
对于单个值的简单情况，这很容易做到（但仍然有重复项）：

a = np.array([1,2,3,4])
mesh = np.array(np.meshgrid(a,a)).T.reshape(-1,2)
>>> out: [[1 1]
          [1 2]
          [1 3]
          [1 4]
          [2 1]
          [...]
          [4 4]]

但对我来说，我的尝试

mesh = np.array(np.meshgrid(ab,ab)).T

给我

[[[ 1  1]
  [ 1 10]
  [ 1  2]
  [ 1 20]
  [ 1  3]
  [ 1 30]
  [ 1  4]
  [ 1 40]]

 [[10  1]
  [10 10]
  [10  2]
  [10 20]
...
  [40  3]
  [40 30]
  [40  4]
  [40 40]]]

换句话说：meshgrid分解了我的配对。我以为解决方案就在附近，但我无法独自提出。任何帮助表示赞赏，谢谢！

最佳答案

不要以为meshgrid在创建所有可能的组合时会起作用（以后不要过滤掉）。为了解决这个问题，可以提出两种方法。

方法1

我们可以获取那些没有重复的成对组合的行索引，然后简单地索引到行中以获得所需的输出，就像这样-

In [99]: r,c = np.triu_indices(len(ab),1)

In [100]: np.hstack(( ab[r], ab[c] ))
Out[100]:
array([[ 1, 10,  2, 20],
       [ 1, 10,  3, 30],
       [ 1, 10,  4, 40],
       [ 2, 20,  3, 30],
       [ 2, 20,  4, 40],
       [ 3, 30,  4, 40]])

要以3D数组的形式获得所需的输出，请沿第二个轴堆叠-

In [115]: np.stack(( ab[r], ab[c] ), axis=1)
Out[115]:
array([[[ 1, 10],
        [ 2, 20]],

       [[ 1, 10],
        [ 3, 30]],

       [[ 1, 10],
        [ 4, 40]],

       [[ 2, 20],
        [ 3, 30]],

       [[ 2, 20],
        [ 4, 40]],

       [[ 3, 30],
        [ 4, 40]]])

作为功​​能：

def pairwise_combs1(ab):
    r,c = np.triu_indices(len(ab),1)
    return np.stack(( ab[r], ab[c] ), axis=1)

方法＃2另一种方法slicing和array-initialization针对内存效率和性能-

def pairwise_combs2(ab):
    n = len(ab)
    N = n*(n-1)//2
    out = np.empty((N,2,2),dtype=ab.dtype)
    idx = np.concatenate(( [0], np.arange(n-1,0,-1).cumsum() ))
    start, stop = idx[:-1], idx[1:]
    for j,i in enumerate(range(n-1)):
        out[start[j]:stop[j],0] = ab[j]
        out[start[j]:stop[j],1] = ab[j+1:]
    return out

运行时测试

In [166]: ab = np.random.randint(0,9,(1000,2))

In [167]: %timeit pairwise_combs1(ab)
10 loops, best of 3: 20 ms per loop

In [168]: %timeit pairwise_combs2(ab)
100 loops, best of 3: 6.25 ms per loop

In [169]: np.allclose(pairwise_combs1(ab), pairwise_combs2(ab))
Out[169]: True