for循环中切片Numpy数组遇索引错误，求批量分组均值解法

我有一个多维numpy弹性（elasticity）数组，维度分别为收入组（低收入/高收入）和年龄组（0-92岁），想通过for循环生成各子组组合的均值表格。代码如下：

import numpy as np

elasticity = np.random.rand(2,92)
print(elasticity.shape)

income = ['i0','i1']
age_gr= [':18','18:']

table = {}
for i in range(len(age_gr)):
    for j in range(len(income)):
        key = age_gr[i]+"_"+income[j]
        table[key] = np.mean(elasticity[age_gr[i],j])
print(table)

运行时触发索引错误：

IndexError: only integers, slices (:), ellipsis (...), numpy.newaxis (None) and integer or boolean arrays are valid indices

实际场景中有更多年龄组，无法手动处理，期望得到各子组的均值结果。补充：年龄维度是0-92岁，并非仅18岁上下划分。

错误原因 #

age_gr里的字符串（比如:18）不能直接作为numpy数组的索引，numpy只接受整数、slice对象、布尔数组等合法索引类型，字符串会被识别为无效索引。

修正代码 #

先写一个工具函数，把年龄组字符串转换成numpy能识别的slice对象，再进行索引计算：

import numpy as np

def str_to_slice(s):
    # 处理':18'、'18:'、'18:30'这类年龄组格式
    parts = s.split(':')
    if len(parts) == 2:
        start = int(parts[0]) if parts[0] else None
        stop = int(parts[1]) if parts[1] else None
        return slice(start, stop)
    raise ValueError("无效的年龄组格式，请使用':n'、'n:'或'n:m'形式")

# 生成模拟数据：收入组(2) × 年龄组(93，对应0-92岁共93个年龄)
elasticity = np.random.rand(2, 93)
income = ['i0', 'i1']
# 示例：扩展更多年龄组划分
age_gr = [':18', '18:30', '30:60', '60:']

table = {}
for age_str in age_gr:
    age_slice = str_to_slice(age_str)
    for idx, inc in enumerate(income):
        key = f"{age_str}_{inc}"
        # 注意数组维度顺序：elasticity[收入组索引, 年龄组slice]
        table[key] = np.mean(elasticity[idx, age_slice])

# 格式化打印结果
for k, v in table.items():
    print(f"{k}: {v:.4f}")

代码说明 #

1. str_to_slice函数：将年龄组字符串解析为numpy可识别的slice对象，支持三种常见划分格式：
   • :n → 匹配0到n-1岁的年龄组（比如:18对应0-17岁）
   • n: → 匹配n到92岁的年龄组（比如18:对应18-92岁）
   • n:m → 匹配n到m-1岁的年龄组（比如18:30对应18-29岁）
2. 循环逻辑：遍历所有年龄组和收入组的组合，用转换后的slice索引数组，计算对应子组的均值并存入字典。
3. 维度注意：原数组elasticity的维度为收入组×年龄组，因此索引顺序为elasticity[收入组索引, 年龄组slice]，和原代码的索引逻辑一致。

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