Векторизация цикла на основе порядка значений в серии

Question

Этот вопрос основан на предыдущем вопросе Я ответил.

Вход выглядит так:

Index   Results  Price
0       Buy      10
1       Sell     11
2       Buy      12
3       Neutral  13
4       Buy      14
5       Sell     15

Мне нужно найти каждую последовательность Buy-Sell (игнорирование дополнительных значений Buy / Sell вне последовательности) и вычисление разницы в Price.

Желаемый результат:

Index Results Price Difference
0     Buy     10    
1     Sell    11    1
2     Buy     12    
3     Neutral 13    
4     Buy     14    
5     Sell    15    3

Мое решение многословно, но, похоже, работает:

from numba import njit

@njit
def get_diffs(results, prices):
    res = np.full(prices.shape, np.nan)
    prev_one, prev_zero = True, False
    for i in range(len(results)):
        if prev_one and (results[i] == 0):
            price_start = prices[i]
            prev_zero, prev_one = True, False
        elif prev_zero and (results[i] == 1):
            res[i] = prices[i] - price_start
            prev_zero, prev_one = False, True
    return res

results = df['Results'].map({'Buy': 0, 'Sell': 1})

df['Difference'] = get_diffs(results.values, df['Price'].values)

Есть ли векторизованный метод?Меня беспокоит удобство сопровождения и производительности кода на большом количестве строк.

---

Редактировать: Код сравнительного анализа:

df = pd.DataFrame.from_dict({'Index': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5},
                             'Results': {0: 'Buy', 1: 'Sell', 2: 'Buy', 3: 'Neutral', 4: 'Buy', 5: 'Sell'},
                             'Price': {0: 10, 1: 11, 2: 12, 3: 13, 4: 14, 5: 15}})

df = pd.concat([df]*10**4, ignore_index=True)

def jpp(df):
    results = df['Results'].map({'Buy': 0, 'Sell': 1})    
    return get_diffs(results.values, df['Price'].values)

%timeit jpp(df)  # 7.99 ms ± 142 µs per loop

Answer 1

Я напишу некоторые альтернативы позже, используя scipy и numpy, но вот четкий, простой ответ только для того, чтобы предложить векторизованную альтернативу, хотя это все равно будет отставать от numba с точки зрения производительности.

Если я правильно понимаю проблему, появится надпись «Купить», за которой следует любое количество возможных альтернатив, и, наконец, появится «Продать», и вы захотите найти разницу между первым «Купить» и «Продать».».Затем начнется другая покупка, и т. Д.

---

Вы можете создать серию для группировки, используя cumsum и shift:

a = df.Results.shift().eq('Sell').cumsum()

0    0
1    0
2    1
3    1
4    1
5    1
Name: Results, dtype: int32

Затем вы можете найти первое и последнее значения для каждой группы, используя agg:

agr = df.groupby(a).Price.agg(['first', 'last'])

Наконец, мы можем назначить новый столбец, используя loc:

df.loc[df.Results.eq('Sell'), 'Diff'] = agr['last'].sub(agr['first']).values

   Index  Results  Price  Diff
0      0      Buy     10   NaN
1      1     Sell     11   1.0
2      2      Buy     12   NaN
3      3  Neutral     13   NaN
4      4      Buy     14   NaN
5      5     Sell     15   3.0

---

Производительность

In [27]: df = pd.concat([df]*10**4, ignore_index=True)

In [28]: %%timeit
    ...: a = df.Results.shift().eq('Sell').cumsum()
    ...: agr = df.groupby(a).Price.agg(['first', 'last'])
    ...: df.loc[df.Results.eq('Sell'), 'Diff'] = agr['last'].sub(agr['first']).values
    ...:
17.6 ms ± 199 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [29]: %%timeit
    ...: s=df.groupby('Results').cumcount()
    ...: df['Diff']=df.Price.groupby(s).diff().loc[df.Results.isin(['Buy','Sell'])]
    ...:
3.71 s ± 331 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Я не могу запустить ваш код, я получаю TypingError так что я не могу сравнить.

Answer 2

Используя cumcount, чтобы найти пару:

s=df.groupby('Results').cumcount()
df['Diff']=df.Price.groupby(s).diff().loc[df.Results.isin(['Buy','Sell'])]
df
Out[596]: 
   Index  Results  Price  Diff
0      0      Buy     10   NaN
1      1     Sell     11   1.0
2      2      Buy     12   NaN
3      3  Neutral     13   NaN
4      4      Buy     14   NaN
5      5     Sell     15   3.0