Question

我有一个矩阵

A = np.zeros((n,n))

假设 n=4。我有另一个矩阵

B = np.array(
    [7, 3, 5, 4],
    [4, 3, 2, 1],
    [6, 7, 4, 5],
    [1, 2, 3, 4]
)

首先我需要 B 的每一行中最小 k 个元素（例如 2）的索引，在这种情况下是

np.array(
    [1, 3],
    [3, 2],
    [2, 3],
    [0, 1]
)

我可以这样做 np.argsort 但最终对整个数组进行排序，np.argpartition 可以对前 k 个元素进行排序，但我需要一种方法来获取索引。最后，我想在给出矩阵的索引处将 B 的最小 k 值插入 A

A = np.array(
    [0, 3, 0, 4],
    [0, 0, 2, 1],
    [0, 0, 4, 5],
    [1, 2, 0, 0]
)

有什么可能的方法？

Answer 1

试试这个：

n = 4
A = np.zeros((n,n), dtype=np.int32)
B = np.array([[7,3,5,4],[4,3,2,1],[6,7,4,5],[1,2,3,4]])

indices = np.argpartition(B, 2, axis=1)[:,:2]
values = np.take_along_axis(B, indices, axis=1)
np.put_along_axis(A, indices, values, axis=1)
print(A)

输出：

array([[0, 3, 0, 4],
       [0, 0, 2, 1],
       [0, 0, 4, 5],
       [1, 2, 0, 0]], dtype=int32)

在特定索引处将一个矩阵中的元素插入另一个矩阵 numpy

1 个答案: