在Python中使用遮罩从矩阵获取列表

时间:2020-08-04 19:15:06

标签: python-3.x numpy

在Python中使用另一个矩阵作为掩码从矩阵中获取列表的最快方法是什么?

列表的顺序必须按列。

示例:

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]

1 个答案:

答案 0 :(得分:1)

请注意,您的matrix/mask看起来像一个元组,对于numpy向量化而言可能并不理想。让我们尝试将它们转换为np.array

matrix = np.array(matrix)

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

get_mask(mask1, matrix)

输出:

array([4, 2, 6])