拼接numpy数组列表时如何同步获取原数组索引及元素位置

你可以通过Numpy内置的批量操作高效生成目标数组，无需手动遍历每个元素，实现逻辑如下：

1. 统计每个子数组的长度
2. 用np.repeat生成列表索引列：每个列表索引重复对应子数组长度次
3. 拼接每个子数组的下标序列生成内部索引列
4. 拼接三列数据得到最终结果

完整实现代码：

import numpy as np

mylist = [
    np.array([48.5, 38.0, 40.0]),
    np.array([61.5, 52.5, 55.5, 46.5]),
    np.array([35.5, 36.5]),
]

# 计算每个子数组的长度
lengths = [len(arr) for arr in mylist]
# 生成列表索引列
list_idx = np.repeat(np.arange(len(mylist)), lengths)
# 生成子数组内部索引列
inner_idx = np.concatenate([np.arange(l) for l in lengths])
# 生成数值列
values = np.concatenate(mylist)
# 拼接得到结果
result = np.vstack((list_idx, inner_idx, values)).T

# 验证结果和期望一致
expected = np.vstack(
    (
        np.array([0, 0, 0, 1, 1, 1, 1, 2, 2]),
        np.array([0, 1, 2, 0, 1, 2, 3, 0, 1]),
        np.array([48.5, 38.0, 40.0, 61.5, 52.5, 55.5, 46.5, 35.5, 36.5]),
    )
).T
print(np.array_equal(result, expected)) # 输出True

如果需要处理超大规模数据，可以用完全向量化的方式生成内部索引，避免多次调用np.arange，性能更高：

lengths = np.array([len(arr) for arr in mylist])
offsets = np.zeros(len(lengths)+1, dtype=int)
offsets[1:] = np.cumsum(lengths)
inner_idx = np.arange(offsets[-1]) - offsets[:-1].repeat(lengths)

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