Python: separate matrix by column values

Question

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.

Answer 1

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

Answer 2

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

Answer 3

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