Python Pandas select rows in numpy array on first columns

Question

I have a dataframe like this:

df = pd.DataFrame({'A':[1,2,3], 'B':[4,5,6],'C':[7,8,9],'D':[10,11,12]})

and a list, here arr, that may vary in length like this:

arr = np.array([[1,4],[2,6]])
arr = np.array([[2,5,8], [1,5,8]])

And I would like get all rows in df that matches first elements in arr like following:

for x in arr:
   df[df.iloc[:, :len(x)].eq(x).all(1)]

Thanks guys!

What do you mean by "get all rows in df that matches first elements in arr" ? What do you define as 'first elements'? What do you define as a match? Please provide an example of the desired output given your example data — itprorh66, Jan 26 '21 at 15:19

score 2 · Accepted Answer · answered Jan 26 '21 at 15:36

IIUC, you can convert the array to df and use merge

arr = np.array([[1,4],[2,6],[2,5]])
df.merge(pd.DataFrame(arr, columns = df.iloc[:,:arr.shape[1]].columns))

    A   B   C   D
0   1   4   7   10
1   2   5   8   11

This solution will handle arrays of different shapes (as long as shape[1] of arr <= shape[1] of df)

arr = np.array([[2,5,8], [1,5,8], [3,6,9]])
df.merge(pd.DataFrame(arr, columns = df.iloc[:,:arr.shape[1]].columns))

    A   B   C   D
0   2   5   8   11
1   3   6   9   12

Python Pandas select rows in numpy array on first columns

1 Answers1