DataFrame拆分后使用iloc取值抛出KeyError:0问题求助

有如下Excel表格：

需求是按POI bus列拆分DataFrame：一个包含前两行（POI bus=0），另一个包含第三行（POI bus≠0）。使用iloc逐行获取指定列数据时，处理第一个DataFrame的gen_sub函数运行正常，输出：

20000
J874
20001
J976

但处理第二个DataFrame的gen_tap函数抛出KeyError:0，报错信息：

in gen_tap
    n4 = bus_num[j]
 
  File ~\AppData\Local\Programs\Spyder\pkgs\pandas\core\series.py:1007 in __getitem__
    return self._get_value(key)
 
  File ~\AppData\Local\Programs\Spyder\pkgs\pandas\core\series.py:1116 in _get_value
    loc = self.index.get_loc(label)
 
  File ~\AppData\Local\Programs\Spyder\pkgs\pandas\core\indexes\base.py:3655 in get_loc
    raise KeyError(key) from err
 
KeyError: 0

使用的代码如下：

from openpyxl import load_workbook
import pandas as pd

input_file = '..\\Bench\\Gen_addition_total.xlsx'

wb = load_workbook(input_file)
ws = wb.active
    
output_file1 = open('..\\Bench\\Gen_sub.inch','w')
output_file2 = open('..\\Bench\\Gen_tap.inch','w')

df0 = pd.read_excel(input_file)
df1 = df0[(df0['POI 2'] == 0)]
df2 = df0[(df0['POI 2'] != 0)]

def gen_sub():
    for i in range(len(df1)):
        bus_num = (df1.iloc[:,0])
        bus_name = (df1.iloc[:,1])
        n1 = bus_num[i]
        n2 = bus_name[i]
        output_file1.write(str(n1).strip("()"))+ output_file1.write('\n') + output_file1.write(str(n2).strip("()")) + output_file1.write('\n')
    output_file1.close()
    
def gen_tap():    
    for j in range(len(df2)):
        bus_num = (df2.iloc[:,0])
        bus_name = (df2.iloc[:,1])
        n4 = bus_num[j]
        n5 = POI_bus2[j]
        output_file2.write(str(n4).strip("()"))+ output_file2.write('\n') + output_file2.write(str(n5).strip("()")) + output_file2.write('\n')
    output_file2.close()
    
gen_sub()  
gen_tap()

1. 索引不连续问题：拆分后的df2保留了原DataFramedf0的索引（第三行索引为2），但代码用range(len(df2))得到循环变量j=0，直接用bus_num[j]去Series里找索引0，而bus_num的实际索引是2，因此抛出KeyError。
2. 未定义变量问题：代码中POI_bus2未定义，属于笔误，应该是要取df2的第三列数据。
3. 低效遍历逻辑：每次循环都重复获取整列数据，既浪费资源又容易引发索引类错误。

方案1：重置DataFrame索引 #

拆分后重置df2的索引，让索引从0开始，同时修复未定义变量问题：

# 修正拆分逻辑，同时重置索引
df1 = df0[(df0['POI 2'] == 0)]
df2 = df0[(df0['POI 2'] != 0)].reset_index(drop=True)

def gen_tap():    
    for j in range(len(df2)):
        bus_num = df2.iloc[:,0]
        poi_bus = df2.iloc[:,2]
        n4 = bus_num[j]
        n5 = poi_bus[j]
        output_file2.write(f"{str(n4).strip('()')}\n{str(n5).strip('()')}\n")
    output_file2.close()

方案2：高效遍历（推荐） #

直接遍历DataFrame的行，彻底避免索引问题，同时优化代码逻辑：

def gen_sub():
    for _, row in df1.iterrows():
        output_file1.write(f"{row.iloc[0]}\n{row.iloc[1]}\n")
    output_file1.close()

def gen_tap():    
    for _, row in df2.iterrows():
        output_file2.write(f"{row.iloc[0]}\n{row.iloc[2]}\n")
    output_file2.close()

方案3：使用列名访问（更健壮） #

如果列名固定，直接用列名访问可避免列顺序变化导致的错误，同时修正代码中列名不一致的问题（原表格列名是POI bus，代码写的POI 2应为笔误）：

# 修正列名匹配问题，同时重置索引
df1 = df0[(df0['POI bus'] == 0)]
df2 = df0[(df0['POI bus'] != 0)].reset_index(drop=True)

def gen_sub():
    for _, row in df1.iterrows():
        output_file1.write(f"{row['Bus num']}\n{row['Bus name']}\n")
    output_file1.close()

def gen_tap():    
    for _, row in df2.iterrows():
        output_file2.write(f"{row['Bus num']}\n{row['POI bus']}\n")
    output_file2.close()

额外优化：用with管理文件 #

避免忘记关闭文件的风险，代码更安全：

def gen_sub():
    with open('..\\Bench\\Gen_sub.inch','w') as f:
        for _, row in df1.iterrows():
            f.write(f"{row['Bus num']}\n{row['Bus name']}\n")

def gen_tap():    
    with open('..\\Bench\\Gen_tap.inch','w') as f:
        for _, row in df2.iterrows():
            f.write(f"{row['Bus num']}\n{row['POI bus']}\n")

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