CODEWITHSUNDEEP
pythonconditional-statementsindices
Forzaa
Forzaa

Reputation: 1545

Python: finding index of an array under several conditions

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.

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.

Upvotes: 3

Views: 15269

Answers (5)

elyase
elyase

Reputation: 41003

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()

Upvotes: 2

Hensing
Hensing

Reputation: 421

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.

Upvotes: 6

unutbu
unutbu

Reputation: 880797

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]

Upvotes: 0

Sukrit Kalra
Sukrit Kalra

Reputation: 34531

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]

Upvotes: 1

njzk2
njzk2

Reputation: 39405

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]

Upvotes: 3

Related Questions