合并两个 NumPy 二维数组

Question

I am trying to merge/join two numpy array based on some column in the most time-efficient way.

Here are my two numpy arrays:

import numpy as np

a1 = np.array([('a', 4.672, 0.999),
   ('b', 2.24, 0.999),
   ('c', 1.984, 0.9001),])

a2 = np.array([('a', 4.67),
   ('c', 2.24)])

Here is my function to return an array of size a1's rows of indices matching on a2's first column . I use this later to make an array of equal size and concatenate. So I need [0, -, 1] for a1 which gives the index of matching rows of a2.

def find_indices(x, y):
    x = x[:, 0]
    y = y[:, 0]

    index = np.argsort(x)
    sorted_x = x[index]
    sorted_index = np.searchsorted(sorted_x, y)

    yindex = np.take(index, sorted_index, mode="clip")
    mask = x[yindex] != y

    result = np.ma.array(yindex, mask=mask)

    return result

find_indices(a2, a1)

This returns this array which seems good.

masked_array(data=[0, --, 1],
             mask=[False,  True, False],
             fill_value=999999)

but when I slice using masked_array.data, I get replicated data,

a2[idx.data]

array([['a', '4.67'],
   ['c', '2.24'],
   ['c', '2.24']], dtype='<U4')

whereas the output should be:

array([['a', '4.67'],
             --     , 
       ['c', '2.24']], dtype='<U4')

How to fix this bug?

Answer 1

Finally, I found the solution:

idx = find_indices(a2, a1)

d = a2[idx]
d[idx.mask] = np.nan

merged_array = np.concatenate((a1, d), axis=1)

I use the mask to replace the duplicated values with Nan. Then directly concatenate.