在公共列上找到兩組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的第一個cols如下匹配；

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

我的方法似乎可以很好地擴展，但是在未排序的大數組上效率不高。 我希望添加到列表中成為主要罪魁禍首。