如何对NumPy数组中1000-2000范围的数值求和并保留列索引？

首先定义原始数组：

import numpy as np

temp = np.array([[   0,    0,    0,    0,    0,    0,    0,    0, 1000,    0,    0,
                    0,    0,    0,    0,    0,    0,    0],
                [   0,    0, 3000, 2000,    0,    0,    0,    0,    0,    0,    0,
                    0,    0,    0,    0,    0,    0,    0],
                [   0,    0,    0,    0,    0,    0,    0, 1500,    0,    0,    0,
                    0,    0,    0,  500,    0,    0,    0]])

实现步骤与代码 #

1. 生成筛选掩码标记目标范围元素
2. 计算每行符合条件的元素和
3. 提取对应列索引

完整实现代码：

# 生成筛选掩码：标记1000到2000（含边界）的元素
mask = (temp >= 1000) & (temp <= 2000)

# 计算每行符合条件的元素和
row_sums = np.where(mask, temp, 0).sum(axis=1)

# 获取每行有效的列索引
col_idx_array = np.arange(temp.shape[1])
valid_col_indices = []
for row_mask in mask:
    valid_cols = col_idx_array[row_mask].tolist()
    valid_col_indices.append(valid_cols)

# 输出结果
for idx in range(len(row_sums)):
    print(f"第{idx+1}行")
    print(f"  符合条件元素的和：{row_sums[idx]}")
    print(f"  对应列索引：{valid_col_indices[idx]}")

代码说明 #

• 掩码mask是与原数组同形状的布尔数组，仅目标范围内元素的位置为True。
• np.where(mask, temp, 0)将不符合条件的元素替换为0，按行求和即可得到每行的目标和。
• 遍历每行掩码，结合列索引数组col_idx_array，直接提取True位置对应的列索引并整理为列表。

运行后输出结果：

第1行
  符合条件元素的和：1000
  对应列索引：[8]
第2行
  符合条件元素的和：2000
  对应列索引：[3]
第3行
  符合条件元素的和：2000
  对应列索引：[7, 14]

内容的提问来源于stack exchange，提问作者Maluma Corazón