根据许多列的极坐标中的另一个列条件（最大值）查找列的值

Question

If I have this dataframe:

pl.DataFrame(dict(x=[0, 1, 2, 3], y=[5, 2, 3, 3],z=[4,7,8,2]))
shape: (4, 3)
┌─────┬─────┬─────┐
│ x   ┆ y   ┆ z   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 0   ┆ 5   ┆ 4   │
│ 1   ┆ 2   ┆ 7   │
│ 2   ┆ 3   ┆ 8   │
│ 3   ┆ 3   ┆ 2   │
└─────┴─────┴─────┘

and I want to find the value in x where y is max, then again find the value in x where z is max, and repeat for hundreds more columns so that I end up with something like:

shape: (2, 2)
┌────────┬─────────┐
│ column ┆ x_value │
│ ---    ┆ ---     │
│ str    ┆ i64     │
╞════════╪═════════╡
│ y      ┆ 0       │
│ z      ┆ 2       │
└────────┴─────────┘

or

shape: (1, 2)
┌─────┬─────┐
│ y   ┆ z   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 0   ┆ 2   │
└─────┴─────┘

What is the best polars way to do that?

Answer 1

You could:

• pl.exclude("x") to select all columns minus x

• .arg_max() to get the max index of each selected column.

• pass the indexes to pl.col("x").take() to get the x value at each index.

• pl.concat_list() to create a list of all the values.

>>> df.select(pl.concat_list(pl.col("x").take(pl.exclude("x").arg_max())))
shape: (1, 1)
┌───────────┐
│ x         │
│ ---       │
│ list[i64] │
╞═══════════╡
│ [0, 2]    │
└───────────┘

To add in the column names you could:

# Feels like this could be simplified?
columns = df.columns
columns.remove("x")
columns = pl.Series(columns).alias("column")

df.select(
   pl.concat_list(
      pl.col("x").take(pl.exclude("x").arg_max())
   ).flatten()
).with_columns(columns)

shape: (2, 2)
┌─────┬────────┐
│ x   | column │
│ --- | ---    │
│ i64 | str    │
╞═════╪════════╡
│ 0   | y      │
├─────┼────────┤
│ 2   | z      │
└─────┴────────┘

Possible approach for the other result:

(df.with_column(pl.exclude("x").arg_max())
   .select([
      pl.col("x").take(col).first().alias(col) 
      for col in df.columns 
      if col != "x"
   ])
)

shape: (1, 2)
┌─────┬─────┐
│ y   | z   │
│ --- | --- │
│ i64 | i64 │
╞═════╪═════╡
│  0  |  2  │
└─────┴─────┘