python - 查找总和为零的向量集

Question

我有 4 个数组(set1、set2、...)，每组 3 个数组。例如

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

这个怎么样？

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 的键。