CODEWITHSUNDEEP
pythonnumpy
Aeolai
Aeolai

Reputation: 199

Extract a subset of data from numpy array

I have a 2D numpy array that I need to extract a subset of data from where the value of the 2nd column is higher than a certain value. What's the best way to do this?

E.g. given the array:

array1 = [[1, 5], [2, 6], [3, 7], [4, 8]]

I would want to extract all rows where the 2nd column was higher than 6, so I'd get:

[3, 7], [4, 8]

Upvotes: 1

Views: 3225

Answers (3)

s3dev
s3dev

Reputation: 9721

Or, even more simply:

a[a[:,1] > 6]

Output:

array([[3, 7], [4, 8]])

Where a is the array.

Upvotes: 2

frab
frab

Reputation: 1173

Use list comprehension:

array1 = [[1, 5], [2, 6], [3, 7], [4, 8]]

threshold = 6
print([elem for elem in array1 if elem[1] > threshold])
# [[3, 7], [4, 8]]

Or using numpy:

import numpy as np

array1 = np.array(array1)
print(array1[array1[:,1] > 6])
# array([[3, 7], [4, 8]])

Upvotes: 1

Oli
Oli

Reputation: 2602

Use numpy.where:

import numpy as np

a = np.array([[1, 5], [2, 6], [3, 7], [4, 8]])
# all elements where the second item it greater than 6:
print(a[np.where(a[:, 1] > 6)])
# output: [[3 7], [4 8]]

Upvotes: 1

Related Questions