如何基于1D数组从两个3D NumPy数组选取元素生成新数组？

你的问题出在维度广播不匹配：c是1维数组（形状(3,))，而a、b是3维数组（形状(3,2,3)）。直接使用np.where(c==0, a, b)时，NumPy会将c广播为(1,1,3)的形状，导致布尔判断作用在数组的最后一维，而非你期望的第一层（轴0），所以结果不符合预期。

下面是三种可行的解决方法：

方法1：扩展c的维度以匹配a、b #

将c扩展为(3,1,1)的形状，让布尔判断作用在第一层（每个"层"）上：

import numpy as np

a = np.array([[[1, 2, 3],
               [4, 5, 6]],

              [[7, 8, 9],
               [10, 11, 12]],

              [[13, 14, 15],
               [16, 17, 18]]])

b = a + 100
c = np.array([0, 1, 0])

# 扩展维度：[:, None, None] 等价于 np.expand_dims(c, (1,2))
result = np.where(c[:, None, None] == 0, a, b)
print(result)

输出结果完全符合预期：

array([[[  1,   2,   3],
        [  4,   5,   6]],

       [[107, 108, 109],
        [110, 111, 112]],

       [[ 13,  14,  15],
        [ 16,  17,  18]]])

方法2：直接按索引循环拼接 #

如果数组规模不大，直接通过循环按c的元素选择对应层，逻辑直观：

result = np.array([a[i] if c[i] == 0 else b[i] for i in range(len(c))])

方法3：堆叠数组后用高级索引选取 #

将a和b堆叠成一个4维数组，再用c作为索引直接提取对应层：

# 沿轴0堆叠a和b，得到形状(2,3,2,3)的数组
stacked = np.stack([a, b], axis=0)
# 用c选择第0轴的元素，np.arange(len(c))匹配第1轴的每个层
result = stacked[c, np.arange(len(c)), :, :]

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