在公共列上查找两组 Python Numpy 数组

Question

我正在尝试合并一些数据，我有以下两个 2d numpy 数组（a 和 b）

a = [[  10   9.689474368e-04][  20   6.88780375e-04]
     [  30   4.296339997e-04][  40  -1.06232578e-03]
     [  50  -1.219884414e-03][  60  -1.27936723e-03]]

b = [[  30   6.687897368e-04][  40   2.887890375e-04]
     [  50   3.293467897e-04][  60  -8.067893578e-03]
     [  70  -7.213988414e-03][  80  -8.278967323e-03]]

我想获得 2 个新的 2d numpy 数组（c 和 d），其中 a 和 b 的第一个列 匹配如下；

c = [[  30   4.296339997e-04][  40  -1.06232578e-03]
     [  50  -1.219884414e-03][  60  -1.27936723e-03]]

d = [[  30   6.687897368e-04][  40   2.887890375e-04]
     [  50   3.293467897e-04][  60  -8.067893578e-03]]

有人知道这样做的优化方法吗？

我尝试过对每个项目进行简单的循环，但是速度不够快，我知道可以通过更优雅的方法来解决它。

我正在使用以下解决方案。它快得多，但不确定它是否是正确的方法。

aHash = map(tuple, a)
bHash = map(tuple, b)

aKey = {x[:1] for x in aHash}
bKey = {x[:1] for x in bHash}

c = np.array([x for x in bHash if x[:1] in aKey])
d = np.array([x for x in aHash if x[:1] in bKey])

谢谢

Answer 1

这是一个我希望很快的解决方案，尤其是在预排序的数据上。

import numpy as np

a = np.array([[  20  ,6.88780375e-04],
     [  30 ,  4.296339997e-04],[  40 , -1.06232578e-03],
     [  50  ,-1.219884414e-03],[  60 , -1.27936723e-03],[  10  ,9.689474368e-04],])

b = np.array([[  30 ,  6.687897368e-04],[  40  , 2.887890375e-04],
     [  50 ,  3.293467897e-04],[  60 , -8.067893578e-03],
     [  70 , -7.213988414e-03],[  80 , -8.278967323e-03],])

a.sort(axis=0)
b.sort(axis=0)

def merge(a, b):
    c = []
    d = []
    ai = 0
    bi = 0
    while(ai < len(a) and bi < len(b)):
        av = a[ai]
        bv = b[bi]
        if av[0] == bv[0]:
            c.append(av)
            d.append(bv)
            ai += 1
            continue
        if av[0] < bv[0]:
            ai += 1
            continue
        else:
            bi += 1
            continue

    return np.array(c), np.array(d)


print merge(a,b)

这是与当前发布的唯一其他方法的比较。这使用了稍微未排序的原始数组（我想对排序方法施加一些惩罚）

Full tests done 100,000 times
while_loop_method = 3.19426544412 sec
hash_map_method = 3.89232874699 sec

这是对大 1000 倍的混洗数组进行的小规模比较。

Full tests done 1,000 times
while_loop_method = 24.1850584226
hash_map_method = 25.9077035996

我的方法似乎可以很好地扩展，但在未排序的大型数组上效率几乎没有。我希望我添加到列表中是罪魁祸首。