查找数组对，例如 array_1 = -array_2

Question

我寻找一种方法从 np.meshgrid(xrange, xrange, xrange) 中找到与 k = -k 相关的所有向量。

目前我这样做:


@numba.njit
def find_pairs(array):
    boolean = np.ones(len(array), dtype=np.bool_)
    pairs = []
    idx = [i for i in range(len(array))]
    while len(idx) > 1:
        e1 = idx[0]
        for e2 in idx:
            if (array[e1] == -array[e2]).all():
                boolean[e2] = False
                pairs.append([e1, e2])
                idx.remove(e1)
                if e2 != e1:
                    idx.remove(e2)
                break 
    return boolean, pairs

# Give array of 3D vectors
krange = np.fft.fftfreq(N)
comb_array = np.array(np.meshgrid(krange, krange, krange)).T.reshape(-1, 3)

# Take idx of the pairs k, -k vector and boolean selection that give position of -k vectors
boolean, pairs = find_pairs(array)

它有效，但执行时间随着 N 的增加而迅速增加...

也许有人已经处理过了？

Answer 1

主要问题是 comb_array 的形状为 (R, 3)，其中 R = N**3 和 find_pairs 中的嵌套循环运行至少在二次方自 idx.remove 运行s 以来的线性时间，并在 for 循环中调用。此外，在某些情况下，for 循环不会更改 idx 的大小，并且循环似乎永远运行（例如 N=4）。

在O(R log R)中解决这个问题的一个方法是对数组进行排序，然后在线性时间内检查相反的值：

import numpy as np
import numba as nb

# Give array of 3D vectors
krange = np.fft.fftfreq(N)
comb_array = np.array(np.meshgrid(krange, krange, krange)).T.reshape(-1, 3)

# Sorting
packed = comb_array.view([('x', 'f8'), ('y', 'f8'), ('z', 'f8')])
idx = np.argsort(packed, axis=0).ravel()
sorted_comb = comb_array[idx]

# Find pairs
@nb.njit
def findPairs(sorted_comb, idx):
    n = idx.size
    boolean = np.zeros(n, dtype=np.bool_)
    pairs = []
    cur = n-1
    for i in range(n):
        while cur >= i:
            if np.all(sorted_comb[i] == -sorted_comb[cur]):
                boolean[idx[i]] = True
                pairs.append([idx[i], idx[cur]])
                cur -= 1
                break
            cur -= 1
    return boolean, pairs

findPairs(sorted_comb, idx)

请注意，该算法假设每一行最多只有一对有效匹配。如果有几个相等的行，则将它们 2 对 2 配对。如果您的目标是在这种情况下提取所有相等行的组合，那么请注意输出将呈指数增长（恕我直言，这是不合理的）。

即使对于 N = 100，此解决方案也非常快。大部分时间都花在效率不高的排序上（不幸的是，Numpy 没有提供一种方法来有效地对行进行字典顺序 argsort，尽管这种操作从根本上来说是昂贵的）。

查找数组对，例如 array_1 = -array_2

Find pairs of array such as array_1 = -array_2

numpy

numpy-ndarray