如何利用另一numpy数组的行值将目标数组每行前N大值转为1其余为0

实现代码 #

import numpy as np

# 初始化输入数组
arr1 = np.array([[4, 1, 3, 2, 5], [5, 2, 4, 1, 3]])
arr2 = np.array([[2], [1]])

# 按行降序排序得到元素对应的索引
arr1_sorted_indices = np.argsort(-arr1, axis=1)

# 生成位置掩码，标记每行前k个需要设为1的位置（k为arr2对应行的数值）
pos = np.arange(arr1.shape[1])[np.newaxis, :]
mask = pos < arr2

# 初始化全零结果数组，按掩码完成赋值
arr1_transformed = np.zeros_like(arr1)
arr1_transformed[np.arange(arr1.shape[0])[:, np.newaxis], arr1_sorted_indices] = mask.astype(int)

print(arr1_transformed)

输出结果 #

[[1 0 0 0 1]
 [1 0 0 0 0]]

逻辑说明 #

• 通过np.argsort(-arr1, axis=1)得到每行元素从大到小排列的索引值
• 生成和arr1列数等长的位置序列，通过和arr2的广播比较，得到每行需要标记为1的前k个位置的掩码
• 利用numpy高级索引特性，将排序后索引对应的位置按掩码赋值为1，其余位置保持默认的0值

内容的提问来源于stack exchange，提问作者Bhanuteja Aryasomayajula