如何获取Polars List中分位数对应的索引？

Question

假设我有以下数据框

df = pl.DataFrame({'x':[[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]]})

要获得第 n 个百分位，我可以执行以下操作：

list_quantile_30 = pl.element().quantile(0.3)
df.select(pl.col('x').list.eval(list_quantile_30))

但是我不知道如何获取百分位数对应的索引？这是我使用 numpy 的方法：

import numpy as np
series = [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]
np.searchsorted(series, np.percentile(series, 30))

有没有一种方法可以在不使用map_elements的情况下以Polars方式做到这一点？

Answer 1

继续您的示例，您可以使用

pl.arg_where

来搜索条件。

df = pl.DataFrame({'x':[[0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]]})

list_quantile_30 = pl.element().quantile(0.3)

df.with_columns(pl.col('x').list.eval(
    pl.arg_where(list_quantile_30 <= pl.element()).first()
).flatten().alias("arg_where"))

shape: (1, 2)
┌────────────────┬───────────┐
│ x              ┆ arg_where │
│ ---            ┆ ---       │
│ list[i64]      ┆ u32       │
╞════════════════╪═══════════╡
│ [0, 2, ... 20] ┆ 3         │
└────────────────┴───────────┘

这说服我在极坐标中也添加

pl.search_sorted

。