如何将(N*J)×M维度的Numpy数组转为标记子数组最大值的布尔数组？

给定维度为$(N*J) \times M$的矩阵$A$，每一列包含$N$个长度为$J$的子向量。需要生成同维度的矩阵$B$：当$A$中对应元素是其所属子向量的最大值时，$B$对应位置为1（假设每个子向量的最大值唯一）。

示例（$N=2$，$J=2$，$M=2$）：

A = 
[1, 0;
 2, 4;
 9, 3;
 7, 8]

B = 
[0, 0;
 1, 1;
 1, 0;
 0, 1]

初步设想结合np.split拆分子向量、取argmax后转换为布尔数组，但需要更清晰的实现方案。

解决方案 #

方法一：维度重塑 + 广播匹配（高效向量操作） #

利用NumPy的维度重塑和广播机制，全程无循环，适合处理大矩阵：

import numpy as np

# 示例矩阵
A = np.array([[1, 0],
              [2, 4],
              [9, 3],
              [7, 8]])
N, J, M = 2, 2, 2

# 1. 将矩阵A重塑为(N, J, M)，拆分出每个子向量
A_reshaped = A.reshape(N, J, M)
# 2. 获取每个子向量的最大值索引（沿J维度）
max_indices = A_reshaped.argmax(axis=1)
# 3. 生成位置索引矩阵，通过广播匹配最大值位置
pos_indices = np.arange(J).reshape(1, J, 1)
B_reshaped = (pos_indices == max_indices.reshape(N, 1, M)).astype(int)
# 4. 重塑回原维度(N*J, M)
B = B_reshaped.reshape(N*J, M)

print(B)
# 输出：
# [[0 0]
#  [1 1]
#  [1 0]
#  [0 1]]

方法二：按列循环处理（逻辑直观） #

如果更习惯按列拆分处理，可采用循环方式，逻辑更易懂：

import numpy as np

A = np.array([[1, 0],
              [2, 4],
              [9, 3],
              [7, 8]])
N, J, M = 2, 2, 2

B = np.zeros_like(A)
for col_idx in range(M):
    # 提取当前列并拆分为N个子向量
    col_subvecs = A[:, col_idx].reshape(N, J)
    # 获取每个子向量的最大值索引
    max_pos_in_subvec = col_subvecs.argmax(axis=1)
    # 计算原矩阵中的行位置
    row_positions = np.arange(N) * J + max_pos_in_subvec
    # 设置对应位置为1
    B[row_positions, col_idx] = 1

print(B)

关键说明 #

• 方法一避免了循环，通过向量操作和广播实现，效率远高于循环，适合大规模数据
• 两种方法均无需使用np.split，直接通过reshape拆分维度更高效
• 因假设最大值唯一，无需处理多最大值的特殊情况

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