如何按索引对列表列进行切片

Question

我有一个熊猫数据帧如下：

df = pd.DataFrame({'A': [[1,2,3],[1,3,5],[4,84,68]]},index=[0,1,0])

我想要的是：

pd.DataFrame({'A': [[1,2],[3,5],[4,84]] },index=[0,1,0])
Out[690]: 
         A
0   [1, 2]
1   [3, 5]
0  [4, 84]

目前的方法是：

[r[i:i+2] for i, r in zip(df.index, df.A)]

Out[695]: [[1, 2], [3, 5], [4, 84]]

然而，当我尝试使用applylambda时，我感到很失落…..并且无法得到预期的输出

到目前为止，我已经尝试了：

df[['A']].apply(lambda x : x[(x.index):(x.index+2)])

你能帮帮我吗，谢谢。你知道吗

Answer 1

a = np.array(df.A.values.tolist())
i = np.arange(len(df))[:, None]
j = np.arange(2) + df.index.values[:, None]

pd.DataFrame(dict(A=a[i, j].tolist()), df.index)

         A
0   [1, 2]
1   [3, 5]
0  [4, 84]