基于位置数组提取Numpy子数组：代码仅生效部分位置的问题排查

我来帮你排查代码问题并实现需求。先理清楚你的场景：你有一组嵌套的numpy数组，需要对应每个子数组，用匹配的位置数组提取指定索引的行，但当前代码逻辑出错，只处理了最后一组位置。

1. 你的数据与期望 #

原始数据 #

import numpy as np
data=[np.array([[0.9,0.6,0.5,0.4,0.7],[0.8,0.0,0.0,0.8,0.2], [0.9,0.0,0.4,0.4,0.3],[0.9,0.6,0.3,0.2,0.5],[0.8,0.0,0.3,0.1,0.5]]), 
      np.array([[0.9,0.0,0.2,0.4,0.3],[0.0,0.2,0.4,0.0,0.0], [0.0,0.0,0.0,0.2,0.0],[0.5,0.0,0.3,0.6,0.8],[0.5,0.6,0.9,0.0,0.0]])]

位置数组 #

positions_non_zero=[np.array([2,3,4]),np.array([1,4])]

期望输出 #

[array([[0.9, 0. , 0.4, 0.4, 0.3], [0.9, 0.6, 0.3, 0.2, 0.5], [0.8, 0. , 0.3, 0.1, 0.5]]), 
 array([[0. , 0.2, 0.4, 0. , 0. ], [0.5, 0.6, 0.9, 0. , 0. ]])]

2. 当前代码的问题 #

你的代码逻辑存在两个关键错误：

• 你在遍历data的每个子数组时，调用的data_target函数会遍历整个positions_non_zero数组，最后只返回最后一次循环的结果（也就是positions_non_zero[1]对应的行），没有把data的子数组和对应的位置数组配对。
• 函数定义放在循环内部是不必要的，而且会导致逻辑混乱。

当前代码：

df_class11=[]
for n in data:
    def data_target(df_class_target):
        for z in df_class_target:
            x_classA=[n[i] for i in z]
            x_classA=np.vstack(x_classA)
        return x_classA
    df_class11.append(data_target(positions_non_zero))
df_class11

3. 修复后的代码 #

我们只需要把data和positions_non_zero一一配对遍历，利用numpy的整数数组索引直接提取行（比列表推导+vstack更简洁高效）：

import numpy as np

data=[np.array([[0.9,0.6,0.5,0.4,0.7],[0.8,0.0,0.0,0.8,0.2], [0.9,0.0,0.4,0.4,0.3],[0.9,0.6,0.3,0.2,0.5],[0.8,0.0,0.3,0.1,0.5]]), 
      np.array([[0.9,0.0,0.2,0.4,0.3],[0.0,0.2,0.4,0.0,0.0], [0.0,0.0,0.0,0.2,0.0],[0.5,0.0,0.3,0.6,0.8],[0.5,0.6,0.9,0.0,0.0]])]
positions_non_zero=[np.array([2,3,4]),np.array([1,4])]

df_class11 = []
# 用zip把data和对应的位置数组配对
for arr, pos in zip(data, positions_non_zero):
    # numpy数组直接用索引数组提取行，不需要vstack
    extracted = arr[pos]
    df_class11.append(extracted)

print(df_class11)

代码解释 #

• zip(data, positions_non_zero)：将data中的每个子数组和positions_non_zero中对应的位置数组一一绑定，确保每个数组只处理自己对应的位置。
• arr[pos]：numpy支持直接用整数数组作为索引，直接提取指定索引的行，返回的就是二维numpy数组，无需额外拼接。

运行这段代码就能得到你想要的期望输出。

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