CODEWITHSUNDEEP
pythonarraysnumpy
user8864088
user8864088

Reputation:

Choosing numpy array such that last element is 1

Suppose we have a dataset like this:

X = 
6  2   1
-2 4  -1
4  1  -1
1  6   1
2  4   1
6  2   1

I would like to get two data from this one having last digit 1 and another having last digit -1.

X0 = 
-2 4  -1
4  1  -1

And, 

X1 = 
6  2   1
1  6   1
2  4   1
6  2   1

How can we do this in numpy efficiently?

In simple python, I could do this like this:

dataset = np.loadtxt('data.txt')

X0, X1 = [], []
for i in range(len(X)):
    if X[i][-1] == 1:
        X0.append(X[i])
    else:
        X1.append(X[i])

This is slow and cumbersome, Numpy is fast and easy so, I would appreciate if there is easier way in numpy. Thanks.

Upvotes: 1

Views: 46

Answers (3)

Pal
Pal

Reputation: 990

import numpy as np
x = np.array(x)
x0 = x[np.where(a[:,2]==-1)]
x1 = x[np.where(a[:,2]==1)]

Upvotes: 0

Stefan Falk
Stefan Falk

Reputation: 25387

You could just use numpy and use slicing to access your data e.g.:

X[X[:, 2] == 1]  # Returns all rows where the third column equals 1

or as a complete example:

import numpy as np

# Random data set
X = np.zeros((6, 3))
X[:3, 2] = 1
X[3:, 2] = -1
np.random.shuffle(X)

print(X[X[:, 2] == 1])
print('-')
print(X[X[:, 2] == -1])

Upvotes: 0

juanpa.arrivillaga
juanpa.arrivillaga

Reputation: 95908

Suppose you have an array:

>>> arr
array([[ 6,  2,  1],
       [-2,  4, -1],
       [ 4,  1, -1],
       [ 1,  6,  1],
       [ 2,  4,  1],
       [ 6,  2,  1]])

Then simply:

>>> mask1 = arr[:, -1] == 1
>>> mask2 = arr[:, -1] == -1
>>> X1 = arr[mask1]
>>> X2 = arr[mask2]

Results:

>>> X1
array([[6, 2, 1],
       [1, 6, 1],
       [2, 4, 1],
       [6, 2, 1]])
>>> X2
array([[-2,  4, -1],
       [ 4,  1, -1]])

Upvotes: 1

Related Questions