查找总和为零的向量集

Question

我有 3 个数组的 4 个数组（set1、set2、...）。例如

set1 = [array([1, 0, 0]), array([-1, 0, 0]), array([0, 1, 0]), ...]

我需要找出有多少向量组合的总和为零。解决这个问题的简单方法是：

for b1 in set1:
  for b2 in set2:
    for b3 in set3:
      for b4 in set4:
        if all(b1 + b2 + b3 + b4 == 0):
          count = count + 1

但是，这需要 O(n^4)，并且基于 3sum 算法，我假设我可以做到 O(n^3)，并且速度非常重要。关于如何在 python 中快速做到这一点的任何线索？

Answer 1

假设输入是一维数组的列表，如问题中提供的示例数据中所列，您似乎可以在对输入列表进行行堆叠后使用broadcasting，就像这样 -

import numpy as np

s1 = np.row_stack((set1))
s2 = np.row_stack((set2))
s3 = np.row_stack((set3))
s4 = np.row_stack((set4))

sums = s4[None,None,None,:,:] + s3[None,None,:,None,:] + s2[None,:,None,None,:] + s1[:,None,None,None,:]
count = (sums.reshape(-1,s1.shape[1])==0).all(1).sum()

示例运行 -

In [319]: set1 = [np.array([1, 0, 0]), np.array([-1, 0, 0]), np.array([0, 1, 0])]
     ...: set2 = [np.array([-1, 0, 0]), np.array([-1, 1, 0])]
     ...: set3 = [np.array([1, 0, 0]), np.array([-1, 0, 0]), np.array([0, 1, 0])]
     ...: set4 = [np.array([1, 0, 0]), np.array([-1, 0, 0]), np.array([0, 1, 0]), np.array([0, 1, 0])]
     ...: 

In [320]: count = 0
     ...: for b1 in set1:
     ...:   for b2 in set2:
     ...:     for b3 in set3:
     ...:       for b4 in set4:
     ...:         if all(b1 + b2 + b3 + b4 == 0):
     ...:           count = count + 1
     ...:           

In [321]: count
Out[321]: 3

In [322]: s1 = np.row_stack((set1))
     ...: s2 = np.row_stack((set2))
     ...: s3 = np.row_stack((set3))
     ...: s4 = np.row_stack((set4))
     ...: 
     ...: sums = s4[None,None,None,:,:] + s3[None,None,:,None,:] + s2[None,:,None,None,:] + s1[:,None,None,None,:]
     ...: count2 = (sums.reshape(-1,s1.shape[1])==0).all(1).sum()
     ...: 

In [323]: count2
Out[323]: 3

Answer 2

它不会改变实际的时间复杂度，但您可以通过告诉 Python 使用例如将其编译为 C 代码来将这些循环加快几百倍。赛通：http://cython.org/。你也可以并行化，因为你写的代码是并行的。一个好的 C 编译器会自动利用这一点，但 Python 不会。

一个实现多更好时间复杂度（O[n^2 log N]）的算法在此处概述：https://cs.stackexchange.com/questions/2973/generalised-3sum-k-sum-problem。我可能会在 Python 中实现所描述的算法并将 Cython 放在它周围。

编辑：

对于展平的数组，您还可以按照以下方式进行您所勾画的操作：

sum2 = np.add.outer(A,B)
sum3 = np.add.outer(sum2,C)
sum4 = np.add.outer(sum3,D)

sum4[i,j,k,l] 现在是 A[i]+B[j]+C[k]+D[l]。零条目的数量是

len(sum4) - np.count_nonzero(sum4)

Answer 3

使用numpy的meshgrid函数：

http://docs.scipy.org/doc/numpy/reference/generated/numpy.meshgrid.html

您需要将初始集合重塑为 1-D，但这并没有损失。

set1 = set1.flatten() // etc

然后调用meshgrid()。它将为您提供 4 个 4-D 数组，每个数组一个。然后只需添加：

a,b,c,d = np.meshgrid(set1, set2, set3, set4)
total = a+b+c+d 

最后，统计整个数组中0的个数：

count = len(total) - np.count_nonzero(sum)

Answer 4

这个怎么样？

from numpy import array
def createset(min, max):
    xr = lambda: xrange(min, max)
    return [ array([x, y, z]) for x in xr() for y in xr() for z in xr() ]

set1 = createset(-3, 3)
set2 = createset(-2, 1)
set3 = createset(-4, 5)
set4 = createset(0, 2)

lookup = {}
for x in set1:
    for y in set2:
        key = tuple(x + y)
        if key not in lookup:
            lookup[key] = 0
        lookup[key] += 1

count = 0
for x in set3:
    for y in set4:
        key = tuple(-1 * (x + y))
        if key in lookup:
            count += lookup[key]

print count

这个想法是生成前两组的所有总和。然后，生成最后两组的总和，并查看查找表中是否有一个键，使得它们的总和为 0。

Answer 5

您可以在sum 函数中使用itertools.product 和生成器表达式：

from itertools import combinations
sum(1 for i in produt(set1,set2,set3,set4) if sum(i)==0)

这将比您的代码更快，但仍然是 O(n⁴) 以获得更快的速度，您可以 get the product with Numpy 而不是 itertools。