如何从多行列数组按条件提取元素并创建新数组？含列求和需求

嘿，你的现有代码已经完美实现了需求——根据第7列（索引6）的0/1值，对第4列（索引4）的数值分别求和，逻辑完全没问题！不过我们可以用Python的特性让代码更简洁高效，给你两种优化方案：

方案1：用生成器表达式简化循环 #

如果你用的是普通的列表型数组，不需要额外依赖库，直接用内置的sum()函数搭配生成器表达式就能一行搞定一个求和，代码更紧凑，内存效率也更高（不需要临时创建列表）：

sumNSal = sum(int(row[4]) for row in dataArray if row[6] == 1)
sumNSalR = sum(int(row[4]) for row in dataArray if row[6] == 0)
print(sumNSal, sumNSalR)

这种写法和你的循环逻辑完全一致，但去掉了冗余的循环结构，可读性也更强。

方案2：用numpy实现向量化加速（适合大数据量） #

如果你的数组是数值型且数据量较大（比如后续可能超过52行），用numpy的向量化操作会比Python原生循环快得多。先把列表转成numpy数组，然后通过布尔索引直接筛选对应行并求和：

import numpy as np

# 将你的列表数组转为numpy数组
data_np = np.array(dataArray)
# 筛选第7列为1的行，取第4列转成int后求和
sumNSal = data_np[data_np[:, 6] == 1, 4].astype(int).sum()
# 筛选第7列为0的行，取第4列转成int后求和
sumNSalR = data_np[data_np[:, 6] == 0, 4].astype(int).sum()

print(sumNSal, sumNSalR)

补充说明 #

两种方案都能得到和你原始代码一样的结果，你可以根据自己的场景选择：

• 小数据量或者不想引入依赖：选方案1
• 大数据量、需要高性能计算：选方案2

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