Применять пользовательскую функцию итеративно в подмножестве строк

Question

Я пытаюсь написать функцию, которая позволяет итеративно выполнять некоторую арифметику для подмножества строк, когда условие выполняется в другом столбце.Мой DataFrame выглядит следующим образом:

        Value                                store flag
0       16051.249                               0    0
36      16140.792                        0.019822    0
0       16150.500                            AAA     1
37      16155.223                         1.24698    0
1       16199.700                            BBB     1
38      16235.732                         1.90162    0
41      16252.594                         2.15627    0
2       16256.300                             CCC    1
42      16260.678                         2.15627    0
1048    17071.513                         14.7752    0
3       17071.600                            DDD     1
1049    17072.347                         14.7752    0
1391    17134.538                         16.7026    0
4       17134.600                            EEE     1
1392    17134.635                         16.7026    0
1675    17227.600                         19.4348    0
5       17227.800                            EFG     1
1676    17228.796                         19.4348    0
1722    17262.189                         20.5822    0
6       17264.300                            XYZ     1
1723    17266.625                         20.6702    0
2630    17442.770                         32.7927    0
7       17442.800                             ZZZ    1
2631    17442.951                         32.7927    0
3068    17517.492                         37.6485    0
8       17517.500                            TTT     1
3069    17518.296                         37.6485    0
3295    17565.776                         38.2871    0
9       17565.800                             SDF    1
3296    17565.888                         38.2871    0
...           ...                             ...  ...

Я хотел бы применить следующую функцию ко всем строкам, где значение flag равно 1:

def f(x):
    return df.iloc[0,1]+(df.iloc[2,1]-df.iloc[0,1])*((df.iloc[1,0]-df.iloc[0,0])/(df.iloc[2,0]-df.iloc[0,0]))

инаконец, поместите возвращаемое значение в словарь с соответствующим значением ключа;например {AAA: 123, BBB:456,...}.

Эта функция требует строки выше и ниже строки, где flag=="1"

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

idx = (df['flag'] == "1").fillna(False)
idx |= idx.shift(1) | idx.shift(2)
idx |= idx.shift(-1) | idx.shift(-2)
df=df[idx]
df.rolling(window=3, min_periods=1).apply(f)[::3].reset_index(drop=True)

, но это не работает!

Поскольку функция зависит от местоположения, я не уверен, как применить ее ко всем триплетам строк, где значение флага1. Любое предложение высоко ценится!

Answer 1

просто сохраняйте состояние и применяйте:

zero_vals = []

def func(row):
    if row.flag == 0:
        zero_vals.append(row)
    elif row.flag == 1:
        # do math here using previous rows of data and current row
        zero_vals.clear()
    else:
        raise ValueError('unexpected flag value')

тогда это просто:

df.apply(func, axis=1)

Answer 2

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

# convert to numeric so that the column can be used for arithmetic calculations
df['store2'] = pd.to_numeric(df.store, errors='coerce')

# calculate the f(x) based on 'Value' and 'store2' column
df['result'] = df.store2.shift(1) + (df.store2.shift(-1) - df.store2.shift(1))*(df.Value - df.Value.shift(1))/(df.Value.shift(-1) - df.Value.shift(1))

# export the resultset:
df.loc[df.flag==1,['store','result']].set_index('store')['result'].to_json()