Панды - найдите некоторые конкретные строки и выведите отфильтрованный кадр данных в новый файл CSV

Question

У меня есть CSV-файл с 15 столбцами и около 17000 строк. Моя проблема состоит в том, чтобы найти в определенном столбце (например: столбец «имя») входную строку, если она совпадает, выведите строку [i], содержащую строку, предыдущую строку [i-1] и следующую строка [i + 1], в порядке i-1, i, i + 1. Повторите процесс до последнего элемента столбца (мой файл данных отформатирован так, что он не содержит дубликатов).

Я использую эту ссылку , чтобы найти строки, и программа работает хорошо. Ниже мой код Python:

import pandas as pd
x = input('Please input the name: ')
df = pd.read_csv("input.csv", sep = ",")

idx = df[df.name.str.contains(x, na=False)].index.tolist()

for i in idx:
    print(df.iloc[[i-1, i, i+1]])

Я хотел бы спросить, как экспортировать отфильтрованные данные выше в новый фрейм данных и вывести их в новый CSV-файл? Я следую за этим ссылка :

df.iloc[[i-1, i, i+1]].to_csv('result.csv', index=True, mode='a')

Выходной файл в порядке, но он не включает имена столбцов, и я также думаю, что он не настолько формален и оптимален в соответствии с автором темы.

Большое спасибо.

Answer 1

Я думаю, вам нужно min и max, чтобы избежать выбора несуществующей строки до первой и после последней совпавшей строки, затем для нового файла сначала сохраните только имена столбцов, а затем в цикле сохраните только данные без заголовка:

df = pd.DataFrame({
        'A':list('abcdef'),
         'B':[4,5,4,5,5,4],
         'C':[7,8,9,4,2,3],
         'D':[1,3,5,7,1,0],
         'E':[5,3,6,9,2,4],
         'name':list('aaabbb')
})

print (df)

#tested matching first row
x = 'a'
#tested matching last row
#x = 'b'

idx = df[df.name.str.contains(x, na=False)].index.tolist()

pd.DataFrame(columns=df.columns).to_csv('result.csv')
for i in idx:
    df1 = df.iloc[[max(0, i-1), i, min(df.index[-1], i+1)]]

    df1.to_csv('result.csv', index=False, mode='a', header=None)
    #if need index values
    #df1.to_csv('result.csv', mode='a', header=None)

Другим решением является использование concat списка DataFrames, а затем сохранение в CSV без режима append:

x = 'a'
idx = df[df.name.str.contains(x, na=False)].index.tolist()

dfs = []
for i in idx:
    dfs.append(df.iloc[[max(0, i-1), i, min(df.index[-1], i+1)]])

#list comprehension alternative
#dfs = [df.iloc[[max(0, i-1), i, min(df.index[-1], i+1)]] for i in idx]

---

pd.concat(dfs).to_csv('result.csv', index=False)
#if need index
#pd.concat(dfs).to_csv('result.csv')

Answer 2

Вы можете попробовать это:

tmp = pd.DataFrame() # A blank dataframe
for i in idx:
    tmp.append(df.iloc[[i-1, i, i+1]]) # Keep appending the rows to the tmp df

Затем присвойте имена столбцов tmp

tmp.columns = df.columns

Напишите это в csv:

tmp.to_csv('result.csv', index=True)

Дайте мне знать, если это поможет.