python - 修改矩阵行和列的每两个最大元素

Question

在 python 中，我有一个矩阵，我想在每一行和每一列中找到两个最大的元素，并将它们的值更改为 1(分别地，我的意思是得到两个矩阵，其中一个修改了行，另一个修改了列)。主要目标是得到一个相应的矩阵，除了我在每行和每列的第 2 个最大元素中放置的那些矩阵之外，其他地方都为零(使用 np.where(mat == 1, 1, 0)。我正在尝试使用 np.argpartition 来这样做，但没有成功。请帮忙。见下图。

Answer 1

这是 np.argpartition 的一种方法-

idx_row = np.argpartition(-a,2,axis=1)[:,:2]
out_row = np.zeros(a.shape,dtype=int)
out_row[np.arange(idx_row.shape[0])[:,None],idx_row] = 1

idx_col = np.argpartition(-a,2,axis=0)[:2]
out_col = np.zeros(a.shape,dtype=int)
out_col[idx_col,np.arange(idx_col.shape[1])] = 1

示例输入、输出-

In [40]: a
Out[40]: 
array([[ 3,  7,  1, -5, 14,  2,  8],
       [ 5,  8,  1,  4, -3,  3, 10],
       [11,  3,  5,  1,  9,  2,  5],
       [ 6,  4, 12,  6,  1, 15,  4],
       [ 8,  2,  0,  1, -2,  3,  5]])

In [41]: out_row
Out[41]: 
array([[0, 0, 0, 0, 1, 0, 1],
       [0, 1, 0, 0, 0, 0, 1],
       [1, 0, 0, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1, 0],
       [1, 0, 0, 0, 0, 0, 1]])

In [42]: out_col
Out[42]: 
array([[0, 1, 0, 0, 1, 0, 1],
       [0, 1, 0, 1, 0, 1, 1],
       [1, 0, 1, 0, 1, 0, 0],
       [0, 0, 1, 1, 0, 1, 0],
       [1, 0, 0, 0, 0, 0, 0]])

---

或者，如果您喜欢紧凑的代码，我们可以跳过初始化并使用 broadcasting直接从 idx_row 和 idx_col 获取输出，就像这样 -

out_row = (idx_row[...,None] == np.arange(a.shape[1])).any(1).astype(int)
out_col = (idx_col[...,None] == np.arange(a.shape[0])).any(0).astype(int).T