У меня есть один DataFrame:
import pandas as pd
df = pd.DataFrame({'col_A':[1,0,3], 'col_B':[2,0,1]})
Мне нужно найти разницу столбцов для строк, где сумма элементов в строке больше 0.
null_index = df[df.sum(axis=1)==0].index
df['col_B'] = (df.loc[~df.index.isin(null_index),'col_B']-df.loc[~df.index.isin(null_index),'col_A'])/df.loc[~df.index.isin(null_inde),'col_A']
Я получаю DataFrame с формой (2,1), но мне нужно (3,1), но для 2-й строки ничего не произойдет
Вы можете использовать df.mask
для маскировки col_B
где сумма строк больше 0.
df.col_B.mask(df.sum(1) > 0, df.col_B - df.col_A, inplace=True)
df
# col_A col_B
# 0 1 1
# 1 0 0
# 2 3 -2
Или делать то же самое, если df.mask
с использованием индексной информации, скорее всего, будет медленнее.
df.loc[df.sum(1) > 0, 'col_B'] = df.col_B - df.col_A
Если вы хотите переписать col_B
с разницей только в строках с суммой больше 0:
mask = df.sum(axis=1) == 0
df.loc[mask, 'col_B'] = df.loc[mask].diff(axis=1)['col_B'].astype(int)
df
col_A col_B
0 1 1
1 0 0
2 3 -2