Search code examples
pythonconditional-statementsindices

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.

  • 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.


Solution

  • 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]