CODEWITHSUNDEEP
python-3.xnumpy
rspena
rspena

Reputation: 63

Get list from matrix using a mask in Python

What's the fastest way to get a list from a matrix using another matrix as a mask in Python?

The order of the list must be by column.

Example:

matrix = ([1, 2, 3],
          [4, 5, 6],
          [7, 8, 9])

mask1 = ([0, 1, 0],
        [1, 0, 1],
        [0, 0, 0])

mask2 = ([0, 0, 0],
        [1, 1, 1],
        [0, 0, 0])

mask3 = ([0, 0, 1],
        [0, 1, 0],
        [1, 0, 0])

output1 = [4, 2, 6]

output2 = [4, 5, 6]

output3 = [7, 5, 3]

Upvotes: 0

Views: 43

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

Note that your matrix/mask looks like a tuple, which might not be ideal for numpy vectorization. Let's try convert them to np.array:

matrix = np.array(matrix)

def get_mask(mask, matrix=matrix):
    # .T so we can get the output ordered by columns
    return matrix.T[np.array(mask, dtype=bool).T]

get_mask(mask1, matrix)

Output:

array([4, 2, 6])

Upvotes: 1

Related Questions