CODEWITHSUNDEEP
pythonnumpymatrixsplit
userk
userk

Reputation: 951

Python: separate matrix by column values

I have a matrix A with 3 columns that looks something like, but much larger:

[[10 15 1.0]
 [21 13 1.0]
 [9  14 0.0]
 [14 24 1.0]
 [21 31 0.0]
 ...]

I want to create two separate matrices: one that contains all of the data with the third column=0.0, and another with all the data with the third column=1.0. So essentially splitting the data by the values 0.0 or 1.0 in the third column.

Upvotes: 7

Views: 8806

Answers (4)

Erfan Sharifi
Erfan Sharifi

Reputation: 41

you can use np.hsplit(ary, indices_or_sections)

Upvotes: 0

wim
wim

Reputation: 362517

>>> a
array([[ 10.,  15.,   1.],
       [ 21.,  13.,   1.],
       [  9.,  14.,   0.],
       [ 14.,  24.,   1.],
       [ 21.,  31.,   0.]])
>>> a[np.where(a[:,-1])]
array([[ 10.,  15.,   1.],
       [ 21.,  13.,   1.],
       [ 14.,  24.,   1.]])
>>> a[np.where(~a[:,-1].astype(bool))]
array([[  9.,  14.,   0.],
       [ 21.,  31.,   0.]])

Upvotes: 4

Óscar López
Óscar López

Reputation: 235984

Here's how we'd separate the matrix using list comprehensions, there's no need to import additional libraries. First, one matrix that contains all of the data with the third column is 0.0:

[x for x in matrix if x[2] == 0.0]

And another matrix with all the data that relates to when the third column is 1.0:

[x for x in matrix if x[2] == 1.0]

For example:

matrix = [[10, 15, 1.0],
          [21, 13, 1.0],
          [ 9, 14, 0.0],
          [14, 24, 1.0],
          [21, 31, 0.0]]

[x for x in matrix if x[2] == 0.0]
=> [[ 9, 14, 0.0],
    [21, 31, 0.0]]

[x for x in matrix if x[2] == 1.0]
=> [[10, 15, 1.0],
    [21, 13, 1.0],
    [14, 24, 1.0]]

Upvotes: 1

mdml
mdml

Reputation: 22872

If you're using Numpy, first find the rows where the third column has your desired value, then extract the rows using indexing.

Demo

>>> import numpy
>>> A = numpy.array([[1, 0, 1],
                     [2, 0, 1],
                     [3, 0, 0],
                     [4, 0, 0],
                     [5, 0, 0]])
>>> A1 = A[A[:, 2] == 1, :] # extract all rows with the third column 1
>>> A0 = A[A[:, 2] == 0, :] # extract all rows with the third column 0
>>> A0
array([[3, 0, 0],
       [4, 0, 0],
       [5, 0, 0]])
>>> A1
array([[1, 0, 1],
       [2, 0, 1]])

Upvotes: 8

Related Questions