如何选取DataFrame数组类型列中每行的前N个元素？

问题场景 #

现有一个值为数组的DataFrame，需要将每个数组缩减为前n个元素，示例如下：
原始DataFrame：

import pandas as pd
df = pd.DataFrame({'a': [[1, 2, 3, 4], [2, 2, 3, 3]], 'b': [[6, 7, 8, 9], [6, 6, 9, 9]]})
print(df)
# 输出：
#              a             b
# 0  [1, 2, 3, 4]  [6, 7, 8, 9]
# 1  [2, 2, 3, 3]  [6, 6, 9, 9]

当n=2时，期望输出：

#        a       b
# 0  [1, 2]  [6, 7]
# 1  [2, 2]  [6, 6]

解决方案 #

方法1：使用applymap批量处理所有列

applymap可对DataFrame的每个单元格执行自定义操作，适合全表所有数组列都需要处理的场景：

n = 2
df = df.applymap(lambda x: x[:n])
print(df)

方法2：针对指定列处理

如果只需要处理部分列，可单独选中目标列后使用applymap：

n = 2
target_cols = ['a', 'b']
df[target_cols] = df[target_cols].applymap(lambda x: x[:n])

方法3：列表推导式（大数据量场景更高效）

对于数据量较大的情况，列表推导式的执行效率通常优于applymap，可以逐列处理：

n = 2
for col in df.columns:
    df[col] = [arr[:n] for arr in df[col]]

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