Найти ближайшую следующую строку со значениями, большими или равными текущей строке ⇐ Python

[Anonymous]

Начиная с этого фрейма данных:
import polars as pl

df_1 = pl.DataFrame({
'name': ['Alpha', 'Alpha', 'Alpha', 'Alpha', 'Alpha'],
'index': [0, 3, 4, 7, 9],
'limit': [12, 18, 11, 5, 9],
'price': [10, 15, 12, 8, 11]
})

Мне нужно добавить новый столбец («min_index»), чтобы указать, при каком индексе (больше текущего) цена равна или превышает текущий предел.
В этом примере выше ожидаемый результат:
┌───────┬───────┬───────┬───────┬───────────┐
│ name ┆ index ┆ limit ┆ price ┆ min_index │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═══════╪═══════╪═══════╪═══════╪═══════════╡
│ Alpha ┆ 0 ┆ 12 ┆ 10 ┆ 3 │
│ Alpha ┆ 3 ┆ 18 ┆ 15 ┆ null │
│ Alpha ┆ 4 ┆ 11 ┆ 12 ┆ 9 │
│ Alpha ┆ 7 ┆ 5 ┆ 8 ┆ 9 │
│ Alpha ┆ 9 ┆ 9 ┆ 11 ┆ null │
└───────┴───────┴───────┴───────┴───────────┘

Пояснение результатов столбца «min_index»:

• 1-я строка, где предел равен 12: начиная со 2-й строки, минимальный индекс, цена которого равна или превышает 12, равен 3.
• 2-я строка, где предел равен 18: начиная с 3-й строки, не существует индекса, цена которого равна или выше чем 18.
• 3-я строка, где предел равен 11: начиная с 4-й строки, минимальный индекс, цена которого равна или превышает 11, равен 9.
• 4-я строка, где предел равен 5: начиная с 5-й строки, минимальный индекс, цена которого равна или превышает 5, равен 9.
• 5-я строка, где предел равен 9: как это В последней строке нет другого индекса, цена которого равна или превышает 9.

Мое решение показано ниже, но как можно было бы это сделать в Polars? Мне удалось решить эту проблему за 8 шагов, но я уверен, что есть более эффективный способ.
# Import Polars.
import polars as pl

# Create a sample DataFrame.
df_1 = pl.DataFrame({
'name': ['Alpha', 'Alpha', 'Alpha', 'Alpha', 'Alpha'],
'index': [0, 3, 4, 7, 9],
'limit': [12, 18, 11, 5, 9],
'price': [10, 15, 12, 8, 11]
})

# Group by name, so that we can vertically stack all row's values into a single list.
df_2 = df_1.group_by('name').agg(pl.all())

# Put the lists with the original DataFrame.
df_3 = df_1.join(
other=df_2,
on='name',
suffix='_list'
)

# Explode the dataframe to long format by exploding the given columns.
df_3 = df_3.explode([
'index_list',
'limit_list',
'price_list',
])

# Filter the DataFrame for the condition we want.
df_3 = df_3.filter(
(pl.col('index_list') > pl.col('index')) &
(pl.col('price_list') >= pl.col('limit'))
)

# Get the minimum index over the index column.
df_3 = df_3.with_columns(
pl.col('index_list').min().over('index').alias('min_index')
)

# Select only the relevant columns and drop duplicates.
df_3 = df_3.select(
pl.col(['index', 'min_index'])
).unique()

# Finally join the result.
df_final = df_1.join(
other=df_3,
on='index',
how='left'
)

print(df_final)


Подробнее здесь: https://stackoverflow.com/questions/793 ... urrent-row