Iterate over numpy array

Question

Given the following array:

x = np.array([[0,2,4,5,5.5,6,7],[4,5,6,7,2,3,4]])

Based on that array I need to create another array that skip the rows unit the value in the first column is >5.

So the result should like like this:

([[5.5,6,7],[2,3,4]])

Any hints for a simple (and fast) method for that problem? Thank you very much for your help!

Answer 1

We could use a boolean array as index for filtering.

>>> x[:, x[0] > 5]
array([[ 5.5,  6. ,  7. ],
       [ 2. ,  3. ,  4. ]])

• x[0] selects the first row
• x[0] > 5 creates an array of boolean, checking whether an element is > 5 or not. (This is [False, False, False, False, True, True, True] .)

• When we write some_array[boolean_array] , we only keep the elements in some_array which the corresponding value in boolean_array is True. For instance,

   >>> numpy.array([2, 4, 6, 8])[numpy.array([True, False, False, True])] array([2, 8]) 

• Since we are going to select columns, the boolean array x[0] > 5 should be placed in the second axis. We select the whole first axis with : . Thus the final expression is x[:, x[0] > 5] .

Answer 2

Or the enumerate function:

    res = []
    for i, _ in enumerate(x):
    res.append([])
    for j, val in enumerate(x[i]):
        if j > 5:
            res[i].append(val)