首页 > 解决方案 > 在特定索引处将元素从一个矩阵插入另一个矩阵

问题描述

我有一个矩阵

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

首先,我需要 的每一行中最小 k 个元素(例如 2)的索引B,在这种情况下,它将是

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

有什么可能的方法呢?

标签: pythonarraysnumpy

解决方案


尝试这个:

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)

推荐阅读