CODEWITHSUNDEEP
pythonnumpy
Matias
Matias

Reputation: 621

Iterate over numpy array

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!

Upvotes: 2

Views: 573

Answers (2)

kennytm
kennytm

Reputation: 523724

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

Upvotes: 3

Christoph
Christoph

Reputation: 606

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)

Upvotes: 0

Related Questions