如何用Numpy高效替换二维数组中每个元素为多个值？

给定nm的二维Numpy数组，需按以下规则将每个元素替换为两个值，得到n2m的数组：

• 1 → [1, 0]
• 2 → [0, 1]
• 0 → [0, 0]

输入示例：[[1,0,2],[2,2,1]]，期望输出：[[1,0,0,0,0,1],[0,1,0,1,1,0]]

原for循环实现虽然直观，但频繁调用np.append会不断重新分配内存，在数组规模较大时效率极低。以下是两种Numpy原生的高效向量化实现方案：

方案一：利用索引映射表 #

预定义一个与替换规则对应的映射数组，通过Numpy的索引直接完成批量替换，再重塑数组形状：

import numpy as np

arr = np.array([[1,0,2],[2,2,1]])
# 索引0对应0的替换值，1对应1的替换值，2对应2的替换值
mapping = np.array([[0, 0], [1, 0], [0, 1]])
# 替换后得到(n, m, 2)的三维数组，再重塑为(n, 2*m)
result = mapping[arr].reshape(arr.shape[0], arr.shape[1] * 2)
print(result)

方案二：向量化条件构造 #

分别生成替换后每个位置的第一个值和第二个值，再通过堆叠拼接成目标数组：

import numpy as np

arr = np.array([[1,0,2],[2,2,1]])
# 生成每个元素替换后的第一个值：元素为1时取1，否则0
first_cols = (arr == 1).astype(int)
# 生成每个元素替换后的第二个值：元素为2时取1，否则0
second_cols = (arr == 2).astype(int)
# 按列维度堆叠成三维数组，再重塑为二维目标数组
result = np.stack([first_cols, second_cols], axis=2).reshape(arr.shape[0], -1)
print(result)

两种方案均为完全向量化操作，底层由C实现，相较于原循环方案，在处理大规模数组时性能提升显著。

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