How to group items by difference in percent in Polars? - python-polars

Question

I'd like to group values so that the difference between each group item remains within a certain percentage. E.g. each time when an item is 5% over the first group element it goes to the new group. As a return I need the first group value. Example with 5% threshold where 'a' is given, 'group' and 'groupFirst' must be calculated:
import polars as pl
df = pl.DataFrame({'a': [100, 103, 105, 106, 105, 104, 103, 106, 100, 102],
'group': [0, 0, 1, 1, 1, 1, 1, 1, 2, 2],
'groupFirst': [100, 100, 105, 105, 105, 105, 105, 105, 100, 100]})
print(df)
shape: (10, 3)
┌─────┬───────┬────────────┐
│ a ┆ group ┆ groupFirst │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═══════╪════════════╡
│ 100 ┆ 0 ┆ 100 │
│ 103 ┆ 0 ┆ 100 │
│ 105 ┆ 1 ┆ 105 │
│ 106 ┆ 1 ┆ 105 │
│ ... ┆ ... ┆ ... │
│ 103 ┆ 1 ┆ 105 │
│ 106 ┆ 1 ┆ 105 │
│ 100 ┆ 2 ┆ 100 │
│ 102 ┆ 2 ┆ 100 │
└─────┴───────┴────────────┘

Answer 1

Say you want to reset cummax when your values exceed the value 6. You could do:
(
df.with_columns(
(pl.col("a") >= 6)
.shift(1)
.fill_null(False)
.cumsum()
.alias("group")
).with_columns(
pl.col("a")
.cummax()
.over(pl.col("group"))
.alias("cummax")
)
)
Example:
In [73]: df = pl.DataFrame({'a': [1, 3, 5, 6, 1, 4, 3, 6, 5, 6]})
In [74]: (
...: df.with_columns(
...: (pl.col("a") >= 6)
...: .shift(1)
...: .fill_null(False)
...: .cumsum()
...: .alias("group")
...: ).with_columns(
...: pl.col("a")
...: .cummax()
...: .over(pl.col("group"))
...: .alias("cummax")
...: )
...: )
Out[74]:
shape: (10, 3)
┌─────┬───────┬────────┐
│ a ┆ group ┆ cummax │
│ --- ┆ --- ┆ --- │
│ i64 ┆ u32 ┆ i64 │
╞═════╪═══════╪════════╡
│ 1 ┆ 0 ┆ 1 │
│ 3 ┆ 0 ┆ 3 │
│ 5 ┆ 0 ┆ 5 │
│ 6 ┆ 0 ┆ 6 │
│ ... ┆ ... ┆ ... │
│ 3 ┆ 1 ┆ 4 │
│ 6 ┆ 1 ┆ 6 │
│ 5 ┆ 2 ┆ 5 │
│ 6 ┆ 2 ┆ 6 │
└─────┴───────┴────────┘

Answer 2

I don't think there is a way to use polars expressions to generate the groups since they always rely on the previous group. That being said the groups can be easily generated in O(n) so the penalty for doing that in python should be minor.
import numpy as np
def make_groups(a, threshold=1.05):
a=np.array(a)
outarray=np.empty(len(a), dtype=a.dtype)
outarray[0]=0
curgroup=a[0]
for indx, cur_a in enumerate(a[1:],1):
if cur_a >= threshold * curgroup or cur_a * threshold <= curgroup:
outarray[indx] = outarray[indx-1] + 1
curgroup=cur_a
else:
outarray[indx] = outarray[indx-1]
return pl.Series(outarray)
So now let's take that to our data.
df = pl.DataFrame({'a': [100, 103, 105, 106, 105, 104, 103, 106, 100, 102]})
We just do a map (incidentally, I tried making make_groups into a np.ufunc but couldn't get it to work).
df \
.with_columns(pl.col('a').map(lambda x: make_groups(x, 1.05)).alias('group')) \
.with_columns((pl.col('a').list().over('group').arr.first()).alias('groupFirst'))
shape: (10, 3)
┌─────┬───────┬────────────┐
│ a ┆ group ┆ groupFirst │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═══════╪════════════╡
│ 100 ┆ 0 ┆ 100 │
│ 103 ┆ 0 ┆ 100 │
│ 105 ┆ 1 ┆ 105 │
│ 106 ┆ 1 ┆ 105 │
│ ... ┆ ... ┆ ... │
│ 103 ┆ 1 ┆ 105 │
│ 106 ┆ 1 ┆ 105 │
│ 100 ┆ 2 ┆ 100 │
│ 102 ┆ 2 ┆ 100 │
└─────┴───────┴────────────┘
By the way, if you just want to use the default threshold then you can just do...
df \
.with_columns(pl.col('a').map(make_groups).alias('group')) \
.with_columns((pl.col('a').list().over('group').arr.first()).alias('groupFirst'))