1

I have some numpy arrays like this

p1  = np.array([140,142,145])
p2  = np.array([130,144,147])
p3  = np.array([150,141,147])
p4  = np.array([150,141,148])

I want to compare the first number in p1 with the first number in p2,p3 and p4, etc.

In this instance I want to find if each element is among the lowest two so that the output is

np.array([True,False,True])
np.array([True,False,True])
np.array([False,True,True])
np.array([False,True,False])

1 Answer 1

3

You could use np.argpartition to find the smallest 2 values for each column:

import numpy as np
p1  = np.array([140,142,145])
p2  = np.array([130,144,147])
p3  = np.array([150,141,147])
p4  = np.array([150,141,148])
P = np.row_stack([p1,p2,p3,p4])

result = np.argpartition(P, 2, axis=0) < 2
print(result)

yields

[[ True False  True]
 [ True False  True]
 [False  True False]
 [False  True False]]

np.argpartition(arr, k) partially sorts arr in ascending order. Each group of k elements is smaller than the next group of k elements, but within each group the elements may not be sorted.


Note that the code above always has exactly 2 True values per column. It finds 2 of the lowest values for each column, but may not find all such values. If you wish to find all such values, you could use

In [302]: P <= P[np.argpartition(P, 2, axis=0), np.arange(P.shape[1])][1]
Out[302]: 
array([[ True, False,  True],
       [ True, False,  True],
       [False,  True,  True],
       [False,  True, False]], dtype=bool)

P[np.argpartition(P, 2, axis=0), np.arange(P.shape[1])] returns P in column-sorted order.

In [5]: P[np.argpartition(P, 2, axis=0), np.arange(P.shape[1])]
Out[5]: 
array([[130, 141, 145],
       [140, 141, 147],
       [150, 142, 147],
       [150, 144, 148]])

P[np.argpartition(P, 2, axis=0), np.arange(P.shape[1])][1] selects the 2nd row. These are the 2nd lowest values in each column.

In [6]: P[np.argpartition(P, 2, axis=0), np.arange(P.shape[1])][1]
Out[6]: array([140, 141, 147])

The comparison P <= np.array([140, 141, 147]) is performed by broadcasting the array on the right-hand side from shape (3,) up to shape (4,3) so the comparison can be done element-wise.

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

Comments

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.