Search Numpy array with multiple values

Question

I have numpy 2d array having duplicate values.

I am searching the array like this.

In [104]: import numpy as np

In [105]: array = np.array

In [106]: a = array([[1, 2, 3],
     ...:            [1, 2, 3],
     ...:            [2, 5, 6],
     ...:            [3, 8, 9],
     ...:            [4, 8, 9],
     ...:            [4, 2, 3],
     ...:            [5, 2, 3])

In [107]: num_list = [1, 4, 5]

In [108]: for i in num_list :
     ...:     print(a[np.where(a[:,0] == num_list)])
     ...:
 [[1 2 3]
 [1 2 3]]
[[4 8 9]
 [4 2 3]]
[[5 2 3]]

The input is list having number similar to column 0 values. The end result I want is the resulting rows in any format like array, list or tuple for example

array([[1, 2, 3],
       [1, 2, 3],
       [4, 8, 9],
       [4, 2, 3],
       [5, 2, 3]])

My code works fine but doesn't seem pythonic. Is there any better searching strategy with multiple values?

like a[np.where(a[:,0] == l)] where only one time lookup is done to get all the values.

my real array is large

sorry for long explanation. I posted similar question on code review but I failed to explain it correctly. — Rahul
– Rahul, Commented Jul 5, 2017 at 8:50

Divakar · Accepted Answer · 2017-07-05 10:19:56Z

6

Approach #1 : Using np.in1d -

a[np.in1d(a[:,0], num_list)]

Approach #2 : Using np.searchsorted -

num_arr = np.sort(num_list) # Sort num_list and get as array

# Get indices of occurrences of first column in num_list
idx = np.searchsorted(num_arr, a[:,0])

# Take care of out of bounds cases
idx[idx==len(num_arr)] = 0 

out = a[a[:,0] == num_arr[idx]]

edited Jul 5, 2017 at 10:19

answered Jul 5, 2017 at 8:54

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

8 Comments

Rahul Over a year ago

Sorry for changing the varialble l to num_list.

Rahul Over a year ago

which one is in your opinion better? I am evaluating though for my case

Divakar Over a year ago

@Scripting.FileSystemObject If you have a large number elements in num_list, I would say searchsorted is worth a try.

Rahul Over a year ago

0.60s v 0.32s for in1d v searchsorted

Rahul Over a year ago

Approach #2 is useless for my case. I don't know but it doesn't give all the search.

|

Nils Werner · Accepted Answer · 2017-07-05 08:54:45Z

1

You can do

a[numpy.in1d(a[:, 0], num_list), :]

answered Jul 5, 2017 at 8:54

Nils Werner

37.2k7 gold badges85 silver badges108 bronze badges

3 Comments

Rahul Over a year ago

I am new to numpy syntax. What is last : is for?

Nils Werner Over a year ago

It means "take everything". So a[0, :] is first row, a[:, 0] is first column etc.

Rahul Over a year ago

Ok. I got confused as Diwakar has ommited : in his answer.

Collectives™ on Stack Overflow

Search Numpy array with multiple values

2 Answers 2

8 Comments

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

8 Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related