Filter the rows in a list of tuples using numpy

user_8307221 0
Viewed 391

I am looking for a quicker way to filter out the list of tuples, using numpy and avoiding loops.

A = [(27157, 4),
       (24814, 0),
       (1047, 2),
       (18265, 2),
       (2857, 4),
       (23854, 2),
       (36881, 0)]

Now I have to filter it based on the second element, i.e. 4. If '4' is present form one list 'B', if not form list 'C'.

That is:

B = [(27157, 4),(2857, 4)]
C = [(24814, 0),(1047, 2),(18265, 2),(23854, 2),(36881, 0)]
indexing numpy arrays python
2 Answers

With NumPy, you can use Boolean indexing to return arrays:

mask = A[:, 1] == 4
B = A[mask]
C = A[~mask]

This requires your input to be a NumPy array:

A = np.array([(27157, 4),
              (24814, 0),
              (1047, 2),
              (18265, 2),
              (2857, 4),
              (23854, 2),
              (36881, 0)])
user_92095460

To be fast you must turn you list of tuple in a more efficient data structure. if you want to keep tuples, you can use a structured array :

dt=dtype([('val',int),('key',int)])
B=ndarray(len(A),dt,array(A))

B[B['key']==4] #--> array([(27157, 4), ( 2857, 4)],...
B[B['key']!=4] #--> array([(24814, 0), ( 1047, 2), (18265, 2), (23854, 2), (36881, 0)],...
user_40162850
Related