python - 找到两对总和为相同值的对

Question

我有我使用的随机二维数组

import numpy as np
from itertools import combinations
n = 50
A = np.random.randint(2, size=(n,n))

我想确定矩阵是否有两对总和为同一行向量的行对。我正在寻找一种快速的方法来做到这一点。我目前的方法只是尝试了所有的可能性。

for pair in  combinations(combinations(range(n), 2), 2):
    if (np.array_equal(A[pair[0][0]] + A[pair[0][1]], A[pair[1][0]] + A[pair[1][1]] )):
        print "Pair found", pair

适用于 n = 100 的方法真的很棒。

Answer 1

这是一个纯 numpy 的解决方案；没有太多的计时，但我必须将 n 推到 500 才能看到我的光标在完成前闪烁一次。虽然它是内存密集型的，并且由于更大的 n 的内存需求而会失败。无论哪种方式，我的直觉是找到这样一个向量的几率无论如何都会随着 n 的增大呈几何级数下降。

import numpy as np

n = 100
A = np.random.randint(2, size=(n,n)).astype(np.int8)

def base3(a):
    """
    pack the last axis of an array in base 3
    40 base 3 numbers per uint64
    """
    S = np.array_split(a, a.shape[-1]//40+1, axis=-1)
    R = np.zeros(shape=a.shape[:-1]+(len(S),), dtype = np.uint64)
    for i in xrange(len(S)):
        s = S[i]
        r = R[...,i]
        for j in xrange(s.shape[-1]):
            r *= 3
            r += s[...,j]
    return R

def unique_count(a):
    """returns counts of unique elements"""
    unique, inverse = np.unique(a, return_inverse=True)
    count = np.zeros(len(unique), np.int)
    np.add.at(count, inverse, 1)
    return unique, count

def voidview(arr):
    """view the last axis of an array as a void object. can be used as a faster form of lexsort"""
    return np.ascontiguousarray(arr).view(np.dtype((np.void, arr.dtype.itemsize * arr.shape[-1]))).reshape(arr.shape[:-1])

def has_pairs_of_pairs(A):
    #optional; convert rows to base 3
    A = base3(A)
    #precompute sums over a lower triangular set of all combinations
    rowsums = sum(A[I] for I in np.tril_indices(n,-1))
    #count the number of times each row occurs by sorting
    #note that this is not quite O(n log n), since the cost of handling each row is also a function of n
    unique, count = unique_count(voidview(rowsums))
    #print if any pairs of pairs exist;
    #computing their indices is left as an excercise for the reader
    return np.any(count>1)

from time import clock
t = clock()
for i in xrange(100):
    print has_pairs_of_pairs(A)
print clock()-t

编辑:包括 base-3 包装；现在 n=2000 是可行的，需要大约 2gb 的内存，以及几秒钟的处理

编辑:添加了一些时间； n=100 在我的 i7m 上每次调用只需要 5 毫秒。