Efficient way to get rows of two arrays that have matching values in fraction of their colums

user_10392393 0
Viewed 185

Say we have two 2-dimensional arrays, array1 and array2, or alternatively their panda.DataFrame equivalent. For some of the rows in array1, there can be only one or no row in array2 that has equal values in column 2,3, and 4, which I want to extract. That is:

array1 = array([[1,4,      5, 656,  1385,   434],
                [11,46,    3, 7356, 585,    74],
                [12,45,    5, 656,  135,    134],
                [112,475,  3, 356,  285,    134],
               ])


array2 = array([[15,44,     5, 656, 1385,   434],
                [151,436,   3, 356, 285,     74],
                [521,406,   5, 656, 135,     24],
                [152,445,   54, 56, 635,     134],
                [72,727,    12, 16, 55,      634]])

and what I am looking for is:

array1_filtered = array([[1,4,       5, 656,  1385,     434],
                         [112,475,   3, 356,  285,      134],
                         [12,45,     5, 656,  135,      134],
                                               ])
array2_filtered = array([[15,   44,      5, 656, 1385,     434],
                         [151, 436,      3, 356, 285,      74],
                         [521, 406,      5, 656, 135,      24]
                                                ])

where both filtered arrays match in terms of their column 2,3, and 4. Even better, have the same matching order in their filtered results like above.

I can think of the following solution but turns out to be quite slow:

array1_filtered = []
array2_filtered = []

for row in array1:
    idx = np.where((array2[:, target_columns] == np.array([row[target_columns]])).all(axis=1))[0]
    if idx.size != 0:
        array1_filtered.append(row.tolist())
        array2_filtered.append(array2[idx,:][0].tolist())
array1_filtered  = np.asarray(array1_filtered)
array2_filtered  = np.asarray(array2_filtered)

Does anyone know a quicker way using numpy or pandas.DataFrame?

pandas dataframe numpy python
2 Answers

Let's try pandas:

cols = [2,3,4]

s = pd.DataFrame(np.vstack([array1,array2])).duplicated(cols, keep=False)

array1_filtered = array1[s.iloc[:len(array1)]]
array2_filtered = array2[s.iloc[-len(array2):]]

Update: Another solution that is slightly better that it doesn't count duplicates inside each array:

cols = [2,3,4]

s = (pd.DataFrame(np.vstack([array1,array2]))
       .assign(arr_no=[0]*len(array1) + [1]*len(array2))
    )

duplicates = s.groupby(cols).arr_no.transform('nunique')==2

array1_filtered = array1[duplicates[s.arr_no==0]]
array2_filtered = array2[duplicates[s.arr_no==1]]

Update: For the same order of matching:

cols = [2,3,4]

s = (pd.DataFrame(np.vstack([array1,array2]))
       .assign(arr_no=[0]*len(array1) + [1]*len(array2))
       .sort_values(cols, kind='mergesort')              # here, sort by `cols`
    )

duplicates = s.groupby(cols).arr_no.transform('nunique')==2

# slicing from s instead of from the original arrays
array1_filtered = s.to_numpy()[duplicates&(s.arr_no==0),:-1]
array2_filtered = s.to_numpy()[duplicates&(s.arr_no==1),:-1]

Output:

# array1_filtered
array([[ 112,  475,    3,  356,  285,  134],
       [  12,   45,    5,  656,  135,  134],
       [   1,    4,    5,  656, 1385,  434]])

# array2_filtered
array([[ 151,  436,    3,  356,  285,   74],
       [ 521,  406,    5,  656,  135,   24],
       [  15,   44,    5,  656, 1385,  434]])
user_42384080

This could be done with a merge on columns 2, 3 and 4 :

cols = [2,3,4]
df1 = pd.DataFrame(array1)
df2 = pd.DataFrame(array2)

df_filtered1 = df1.merge(df2.iloc[:,cols], how = 'inner', on = cols).sort_values(cols)
df_filtered2 = df2.merge(df1.iloc[:,cols], how = 'inner', on = cols).sort_values(cols)

# Optional : if you want arrays
array1_filtered = df_filtered1.values
array2_filtered = df_filtered2.values
user_70957610
Related