时间戳

首页 > 解决方案 > 查找 ndarray 与 ndarray 比较的索引

python - 查找 ndarray 与 ndarray 比较的索引

问题描述

我有两个未排序的 ndarray，其结构如下：

a1 = np.array([0,4,2,3],[0,2,5,6],[2,3,7,4],[6,0,9,8],[9,0,6,7])
a2 = np.array([3,4,2],[0,6,9])

我想找到 a1 的所有索引，其中每个 a2 行都在 a1 中，并且在 a1 中的位置：

result = [[0,[3,1,2]],[2,[1,3,0]],[3,[1,0,2]],[4,[1,2,0]]

在此示例中，a2[0] 位于 a1 中的位置 0 和 2 中，位于 3,1,2 和 1,3,0 的 a1 位置。对于 a2[1] 在 1,0,2 和 1,2,0 的 a1 位置内的位置 3 和 4。

每个 a2 行在 a1 中出现两次。a1 至少有 1Mio。行，a2 大约 10,000。所以算法也应该很快（如果可能的话）。

到目前为止，我正在考虑这种方法：

big_res = []
for r in xrange(len(a2)):
    big_indices = np.argwhere(a1 == a2[r])
    small_res = []
    for k in xrange(2):
        small_indices = [i for i in a2[r] if i in a1[big_indices[k]]]
        np.append(small_res, small_indices)
    combined_res = [[big_indices[0],small_res[0]],[big_indices[1],small_res[1]]]
    np.append(big_res, combined_res)

标签： pythonnumpy

解决方案

使用numpy_indexed，（免责声明：我是它的作者）我认为困难的部分可以有效地编写如下：

import numpy_indexed as npi

a1s = np.sort(a1, axis=1)
a2s = np.sort(a2, axis=1)
matches = np.array([npi.indices(a2s, np.delete(a1s, i, axis=1), missing=-1) for i in range(4)])
rows, cols = np.argwhere(matches != -1).T
a1idx = cols
a2idx = matches[rows, cols]
# results.shape = [len(a2), 2]
result = npi.group_by(a2idx).split_array_as_array(a1idx)

这只会给你有效的匹配；不是相对顺序。但是一旦你有了匹配，计算相对顺序应该很容易在线性时间内完成。

编辑：以及一些有问题的密度代码来获得您的相对顺序：

order = npi.indices(
    (np.indices(a1.shape)[0].flatten(), a1.flatten()),
    (np.repeat(result.flatten(), 3),    np.repeat(a2, 2, axis=0).flatten())
).reshape(-1, 2, 3) - result[..., None] * 4

推荐阅读