如何以向量化方式提取DataFrame列表列中对应索引的值？

现有DataFrame #

创建代码：

import pandas as pd

data = {
    'lists': [[0, 1, 2],[3, 4, 5],[6, 7, 8]],
    'indexes': [0, 1, 2]
}
df = pd.DataFrame(data=data)

DataFrame内容：

lists  indexes
0  [0, 1, 2]        0
1  [3, 4, 5]        1
2  [6, 7, 8]        2

需求 #

新增一列extracted_value，存储lists列中每个列表在indexes列对应索引位置的元素，期望结果：

lists  indexes  extracted_value
0  [0, 1, 2]        0                0
1  [3, 4, 5]        1                4
2  [6, 7, 8]        2                8

无效尝试 #

尝试以下代码时，得到的是整个列表而非对应索引的元素：

df['extracted_value'] = df['lists'][df['indexes']]

结果：

lists  indexes extracted_value
0  [0, 1, 2]        0       [0, 1, 2]
1  [3, 4, 5]        1       [3, 4, 5]
2  [6, 7, 8]        2       [6, 7, 8]

同样，以下代码也会得到整个列表：

df['extracted_value'] = df['lists'][0]

高效解决方案 #

针对百万级行的DataFrame，推荐以下两种高效方法：

方法1：Numpy向量化操作（最快） #

利用Numpy的矩阵索引能力，完全避免Python循环，效率最高：

import numpy as np

# 将lists列转换为二维Numpy数组
lists_arr = np.array(df['lists'].tolist())
# 按行索引和对应indexes值提取元素
df['extracted_value'] = lists_arr[np.arange(len(lists_arr)), df['indexes']]

方法2：列表推导式（简洁且高效） #

纯Python列表推导，比apply快数倍，适合数据量较大的场景：

df['extracted_value'] = [lst[idx] for lst, idx in zip(df['lists'], df['indexes'])]

方法3：apply（可读性好，适合小数据量） #

如果对可读性要求更高，数据量不是极端大，也可以用apply：

df['extracted_value'] = df.apply(lambda row: row['lists'][row['indexes']], axis=1)

注意：apply的效率低于前两种方法，百万级数据不推荐。

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