Python中如何合并同维度NumPy数组并按规则替换指定元素？

你的需求可以直接用NumPy原生功能实现，无需引入掩码数组，所有操作对1/2/3维乃至任意维度的同形状数组均通用。

1. 简单场景实现 #

你要求的「保留x2非零值，0值替换为x1对应位置值」可以直接用np.where实现，该函数逐位置判断规则，不会对结果排序，完全符合要求：

import numpy as np
x1 = np.arange(1,10)
x2 = np.array([0,0,4,0,0,5,0,0,0])

# 规则：x2≠0时取x2，否则取x1
res = np.where(x2 != 0, x2, x1)

输出结果：array([1, 2, 4, 4, 5, 5, 7, 8, 9])

如果需要原地修改x2，不需要生成新数组，也可以用布尔索引直接赋值：

x2[x2 == 0] = x1[x2 == 0]

2. 多维场景适配 #

上述两种方法对任意维度的同形状数组都有效，不需要修改逻辑，np.where和布尔索引会自动匹配数组形状逐元素处理。

3. 实际用例&报错解决 #

报错原因 #

之前的代码报错是因为khan、emss函数返回值的维度和索引取出的1维数组不匹配，比如传入1维数组后函数返回了2维结果，赋值时维度对不上导致报错。

最优实现方案 #

你选择的「先分别计算两个函数的结果再合并」的思路效率最高，直接用np.where一行即可完成，会自动处理维度广播适配：

# 前置计算得到同形状的k_angle、e_angle
k_angle = khan(_angle)
e_angle = emss(_angle)

# 规则：_angle<5时取k_angle，否则取e_angle，自动广播匹配维度
final_angle = np.where(_angle < 5, k_angle, e_angle)

如果你还是想用索引赋值的方式，只需要把函数返回结果拉平为1维即可：

test[gt_idx] = khan(test[gt_idx]).ravel()
test[le_idx] = emss(test[le_idx]).ravel()

内容的提问来源于stack exchange，提问作者B.Eng