Сокращение значений столбцов в CSV с использованием сложных операторов If Then

Question

У меня есть CSV-файл со столбцом идентификаторов, столбцами за 6 месяцев, 2 сезонными столбцами и 1 годовым столбцом. Например:

`[['ID', '20180922', '20180820', '20180728', '20180524', '20180424', '20180322', Winter, Summer, Annual], 
['1', '10', '2', '4', '4', '4', '4', '0', '0', '0'],
['2', '4', '2', '2', '10', '10', '4', '0', '0', '0'],
['3', '10', '2', '4', '4', '2', '4', '0', '0', '0'],
['4', '2', '2', '2', '4', '10', '10', '0', '0', '0'],
['5', '10', '4', '4', '2', '10', '10', '0', '0', '0'],
['6', '2', '2', '4', '10', '10', '10', '0', '0', '0'],
['7', '10', '10', '2', '4', '4', '10', '0', '0', '0']]`

Первые три месяца (сентябрь, август, июль) - летние месяцы, а следующие три месяца (май, апрель, март) - зимние месяцы. Значения в таблице являются кодами культур. Я пытаюсь уменьшить значения в сезоны, а затем в год.

Так для Зимней колонны:

• Если в зимние месяцы (май, апрель, март) есть две цифры, заполните столбец Зима цифрой 2
• В противном случае заполните столбец Зима тем, что есть в мае

Для летней колонки:

• Если в летние месяцы есть две цифры (сентябрь, август, июль), заполните колонку «Лето» 2
• В противном случае заполните колонку Summer 10

Для годовой колонки:

• Если в столбце «Зима» или «Лето» есть 2, заполните Годовой 2
• Все остальное 10

Вот как я это делаю сейчас, но я хотел знать, есть ли более чистый способ сделать это. Спасибо

inputFileName = 'input.csv'
outputFileName = 'output.csv'

with open(inputFileName, 'rb') as inFile, open(outputFileName, 'wb') as outfile:
    r = csv.reader(inFile)
    w = csv.writer(outfile)

    table = list(r)

    w.writerow(table[0])

# Winter
for line in table:
    if line[4] == '2' or line[5] == '2' or line[6] == '2':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '2', line[8], line[9]))

    elif line[4] == '4' and line[5] == '4' and line[6] == '10':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '4', line[8], line[9]))

    elif line[4] == '4' and line[5] == '10' and line[6] == '10':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '4', line[8], line[9]))

    elif line[4] == '10' and line[5] == '4' and line[6] == '10':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '10', line[8], line[9]))

    elif line[4] == '10' and line[5] == '10' and line[6] == '4':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '10', line[8], line[9]))

    elif line[4] == '10' and line[5] == '4' and line[6] == '4':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '10', line[8], line[9]))

    elif line[4] == '10' and line[5] == '10' and line[6] == '10':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '10', line[8], line[9]))

    elif line[4] == '4' and line[5] == '4' and line[6] == '4':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], '4', line[8], line[9]))


# Summer
for row in table:
    if line[1] == '2' or line[2] == '2' or line[3] == '2':
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], line[7], '2', line[9]))

    else:
        w.writerow((line[0], line[1], line[2], line[3], line[4], line[5], line[6], line[7], '10', line[9]))

Кроме того, часть кода Summer, похоже, не работает.

Answer 1

Вы можете проверить, есть ли 2 в зимние месяцы с:

if '2' in [line[i] for i in range(4,7)]

Кроме того, у вас есть несколько одинаковых условий, за исключением заказа. Вы можете сделать:

winterlist = [line[i] for i in range(4,7)]
winterlist.sort()

Тогда вы можете просто проверить winterlist, и вам не нужны разные случаи для разных заказов.

А в вашем летнем коде изменение for row in table на for line in table может решить проблему.

Answer 2

Использование панд:

import pandas as pd

data = '''\
ID,20180922,20180820,20180728,20180524,20180424,20180322,Winter,Summer,Annual
1,10,2,4,4,4,4,0,0,0
2,4,2,2,10,10,4,0,0,0
3,10,2,4,4,2,4,0,0,0
4,2,2,2,4,10,10,0,0,0
5,10,4,4,2,10,10,0,0,0
6,2,2,4,10,10,10,0,0,0
7,10,10,2,4,4,10,0,0,0'''

df = pd.read_csv(pd.compat.StringIO(data)).set_index('ID')
#df = pd.read_csv('path/to/file').set_index('ID')

df['Winter'] = df.iloc[:,3:7].apply(lambda x: 2 if 2 in x.values else x[0], axis=1)
df['Summer'] = df.iloc[:,:3].apply(lambda x: 2 if 2 in x.values else 10, axis=1)

cols = ['Winter','Summer']
df['Annual'] = df[cols].apply(lambda x: 2 if 2 in x.values else 10, axis=1)

print(df)

Возвращает:

    20180922  20180820  20180728  20180524  20180424  20180322  Winter  \
ID                                                                       
1         10         2         4         4         4         4       4   
2          4         2         2        10        10         4      10   
3         10         2         4         4         2         4       2   
4          2         2         2         4        10        10       4   
5         10         4         4         2        10        10       2   
6          2         2         4        10        10        10      10   
7         10        10         2         4         4        10       4   

    Summer  Annual  
ID                  
1        2       2  
2        2       2  
3        2       2  
4        2       2  
5       10       2  
6        2       2  
7        2       2