如何用列表推导式遍历Pandas转换的嵌套列表？报错求解

问题描述 #

通过以下代码从Pandas DataFrame生成两个嵌套列表：

next_df = df.shift(-1)
next_waypt = next_df.values.tolist()
waypt=df.values.tolist()

想要用列表推导式结合两个列表的值创建新列表，但遍历索引时出现两个错误：

• 执行以下代码时触发list index out of range错误：

  y = [math.sin(waypt[:1][1] - next_waypt[:1][1])*math.cos(next_waypt[:1][0]) for y in waypt]
  

• 执行另一代码时触发list indices must be integers or slices错误：

  y = [math.sin(waypt[y][1] - next_waypt[y][1])*math.cos(next_waypt[y][0]) for y in waypt]
  

希望找到正确的实现方式，且不想单独提取列以保持代码可读性。

正确实现方式 #

方案1：用enumerate遍历索引（列表推导式版） #

用enumerate同时获取元素的索引和值，就能对应访问waypt和next_waypt的同位置元素。注意shift(-1)后，next_waypt的最后一个元素是NaN（转成列表后可能是None或包含NaN的子列表），需要过滤掉：

import math
import pandas as pd

y = [math.sin(current_pt[1] - next_pt[1]) * math.cos(next_pt[0])
     for idx, current_pt in enumerate(waypt)
     if idx < len(next_waypt) and not any(pd.isna(p) for p in next_pt)]

方案2：直接在Pandas中运算（最优方案） #

没必要转成列表处理，直接在DataFrame上操作更高效、可读性更强：

import pandas as pd
import numpy as np

# 生成下一行的x、y列
df['next_x'] = df.iloc[:, 0].shift(-1)
df['next_y'] = df.iloc[:, 1].shift(-1)

# 矢量化运算（比apply更快）
df['y'] = np.sin(df.iloc[:, 1] - df['next_y']) * np.cos(df['next_x'])

# 转成列表（过滤掉最后一行的NaN）
y_list = df['y'].dropna().tolist()

错误原因解析 #

• 第一个错误：waypt[:1]是取waypt的第一个元素组成的子列表，再取[1]自然越界，因为这个子列表只有索引0的元素。
• 第二个错误：for y in waypt中的y是列表里的子列表元素，不是整数索引，用它去索引waypt就会触发类型错误。

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