python - 使用 Numpy Python 限制矩阵每行中的元素数量

Question

我有以下 numpy 矩阵:

import numpy as np
matrix = np.array([[1,2,3],[4,5,6]])

和一个 numpy 向量:

vector = np.array([1,2])

向量中的每个元素代表矩阵的每一行，我想保留的元素数量。我想用 0 替换矩阵中的所有其他元素。

最终的矩阵应该是这样的:

matrix_output = np.array([[1,0,0],[4,5,0]])

什么是最快的方法？

Answer 1

可以像这样做一些简单的事情:

import numpy as np

matrix = np.array([[1,2,3],[4,5,6]])
vector = np.array([1,2])

for row, index in enumerate(vector):
    matrix[row, index:] = 0

print(matrix)

[[1 0 0]
 [4 5 0]]