Элегантное суммирование столбцов DataFrame в python

Question

Я пытаюсь написать функцию, которая будет суммировать набор указанных столбцов в pandas DataFrame.

Во-первых, некоторые фоне. Каждый из них имеет столбец с именем (например, «var») и рядом с ним имя в последовательном порядке (например, «var1, var2»). Я знаю, что могу подвести, скажем, 5 столбцов вместе со следующим кодом:

import pandas as pd 
data = pd.read_csv('data_file.csv') 
data['var_total'] = data.var1 + data.var2 + data.var3 + data.var4 + data.var5 

Однако, это может быть повторяющиеся, когда у вас есть var1-var30 подводить. Я решил, что должно быть какое-то изящное решение для суммирования их быстрее, поскольку имена столбцов предсказуемы и единообразны. Есть ли функция, которую я могу написать, или встроенную функцию pandas, которая позволит мне суммировать их быстрее?

Answer 1

Даже если вы пишете все имена столбцов есть несколько способов, чтобы сделать сумму чуть более элегантно:

import pandas as pd 
import numpy as np 

df = pd.DataFrame({'var1': np.random.randint(1, 10, 10), 
        'var2': np.random.randint(1, 10, 10), 
        'var3': np.random.randint(1, 10, 10)}) 

# Use the sum method: 
df[['var1', 'var2', 'var3']].sum(axis='columns') 

# Use eval 
df.eval('var1 + var2 + var3') 

Тогда вы всегда можете использовать стандартные инструменты Python для манипулирования строк, чтобы собрать список имен столбцов:

cols = ['var' + str(n) for n in range(1, 3 + 1)] 
cols 
Out[9]: ['var1', 'var2', 'var3'] 

df[cols].sum(axis='columns') 

Answer 2

Я думаю, что вы ищете метод filterDataFrame; вы можете передать ему строку или регулярное выражение, и оно просто вернет столбцы, имена которых соответствуют ему. Тогда вы можете просто позвонить sum или то, что вы хотите на результирующих столбцах:

pd.DataFrame({'var1':[1], 'var2':[2],'othercol':['abc']}) 
    othercol var1 var2 
0  abc  1  2 

pd.DataFrame({'var1':[1], 'var2':[2],'othercol':['abc']}).filter(like='var') 
    var1 var2 
0  1  2 

pd.DataFrame({'var1':[1], 'var2':[2],'othercol':['abc']}).filter(like='var').sum(axis=1) 

0 3 

Кстати заметит, что я назвал sum(axis=1) вернуть строчную сумму, по умолчанию, sum возвращают сумму столбцы.

Answer 3

Вы могли бы сделать что-то вроде этого:

data['var_total'] = data.filter(regex='var[0-9]+').sum(axis=1) 

Сначала будет фильтроваться ваш фреймворк, чтобы сохранить только столбцы, начинающиеся с var, и за ними следует одно или несколько номеров. Затем он будет суммировать полученную фильтрацию DataFrame.