Python: finding index of an array under several conditions

Question

I have the following problem. There are two n-dimensional arrays of integers and I need to determine the index of an item that fulfills several conditions.

• The index should have a negative element in "array1".
• Of this subset with negative elements, it should have the smallest value in "array2".
• In case of a tie, select the value that has the smallest value in "array1" (or the first otherwise)

So suppose we have:

array1 = np.array([1,-1,-2])
array2 = np.array([0,1,1])

Then it should return index 2 (the third number). I'm trying to program this as follows:

import numpy as np
n = 3
array1 = np.array([1,-1,-2])
array2 = np.array([0,1,1])
indices = [i for i in range(n) if array1[i]<0] 
indices2 = [i for i in indices if array2[i] == min(array2[indices])] 
index = [i for i in indices2 if array1[i] == min(array1[indices2])][0] #[0] breaks the tie.

This seems to work, however, I don't find it very elegant. To me it seems like you should be able to do this in one or two lines and with defining less new variables. Anyone got a suggestion for improvement? Thanks in advance.

Answer 1

I don't know much about numpy (though apparently i should really look into it), so here is a plain python solution

This

sorted([(y, x, index) for (index, (x, y)) in enumerate(zip(array1, array2)) if x < 0])

will give you the tripley of elements from array2, array1, index sorted by value in array2 and value in array1 in case of tie, index in case of tie

The first element is what you seek. This gives the following result :

[(1, -2, 2), (1, -1, 1)]

The index is therefore 2, and is obtained by [0][2]

Answer 2

you can get the indices of all negative elements from array1 with:

np.where(array1 < 0)

then you can access the subset via:

subset = array2[array1 < 0]

to get the index of the smallest (negative) value of array1, you can use array1.argmin()

value = array2[array1.argmin()]

putting all together gives you:

value = array2[array1 < 0][array1[array1 < 0].argmin()]

but you have to catch ValueErrors, if array1 has only positive values.

Answer 3

The index should have a negative element in "array1":

indices_1 = np.argwhere(array1 < 0)

Of this subset with negative elements, it should have the smallest value in "array2".

indices_2 = array2[indices_1].argmin()

In case of a tie, select the value that has the smallest value in "array1" (or the first otherwise)

array1[indices_2].argmin()

Answer 4

I came up with this. Just devised one more test case and seemed to work. See if it suits your needs.

sorted([(array2[index], element, index) for index, element in enumerate(array1) if element<0])[0][2]

Or through using the min function.

min([(array2[index], element, index) for index, element in enumerate(array1) if element<0])[2]

Answer 5

The following will work with numpy arrays of any dimension. It uses numpy.lexsort to order the indices.

numpy.lexsort(Y,X) sorts the items in X in ascending order, and breaks ties according to the values in Y. It returns the indices in order (not the values of X or Y.) This is the only function I know of in numpy which "breaks ties" for you.

---

import numpy as np
def find_index(array1, array2):
    indices = np.where(array1 < 0)
    i = np.lexsort((array1[indices],array2[indices]))[0]
    return [idx[i] for idx in indices]

array1 = np.array([1,-1,-2])
array2 = np.array([0,1,1])

array3 = np.array([(1, 2),
                   (-1, -2),
                   (-2, 0)])
array4 = np.array([(0, 2),
                   (1, 1),
                   (3, 0)])

print(find_index(array1, array2))
# [2]

print(find_index(array3, array4))
# [1, 1]