获取 Polars 中的最小/最大列名称

Question

在极坐标中，我可以获得水平最大值（到达行的一组列的最大值），如下所示：

\n

df = pl.DataFrame(\n    {\n        "a": [1, 8, 3],\n        "b": [4, 5, None],\n    }\n)\n\ndf.with_columns(max = pl.max_horizontal("a", "b"))\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 a   \xe2\x94\x86 b    \xe2\x94\x86 max \xe2\x94\x82\n\xe2\x94\x82 --- \xe2\x94\x86 ---  \xe2\x94\x86 --- \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64  \xe2\x94\x86 i64 \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1   \xe2\x94\x86 4    \xe2\x94\x86 4   \xe2\x94\x82\n\xe2\x94\x82 8   \xe2\x94\x86 5    \xe2\x94\x86 8   \xe2\x94\x82\n\xe2\x94\x82 3   \xe2\x94\x86 null \xe2\x94\x86 3   \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

这对应于 Pandas df[["a", "b"]].max(axis=1)。

\n

现在，我如何获取列名称而不是实际的最大值？\n换句话说，Pandas' 的 Polars 版本是什么df[CHANGE_COLS].idxmax(axis=1)？

\n

预期输出为：

\n

\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 a   \xe2\x94\x86 b    \xe2\x94\x86 max \xe2\x94\x82\n\xe2\x94\x82 --- \xe2\x94\x86 ---  \xe2\x94\x86 --- \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64  \xe2\x94\x86 str \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1   \xe2\x94\x86 4    \xe2\x94\x86 b   \xe2\x94\x82\n\xe2\x94\x82 8   \xe2\x94\x86 5    \xe2\x94\x86 a   \xe2\x94\x82\n\xe2\x94\x82 3   \xe2\x94\x86 null \xe2\x94\x86 a   \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

Answer 1

您可以使用 将元素连接到列表中pl.concat_list，使用 获取最大元素的索引pl.Expr.list.arg_max，并使用 将该索引替换为列名pl.Expr.replace。

\n

mapping = {0: "a", 1: "b"}\n(\n    df\n    .with_columns(\n        pl.concat_list(["a", "b"]).list.arg_max().replace(mapping).alias("max_col")\n    )\n)\n

\n

这一切都可以包装到一个函数中来处理映射字典的创建。

\n

def max_col(cols) -> str:\n    mapping = dict(enumerate(cols))\n    return pl.concat_list(cols).list.arg_max().replace(mapping)\n\ndf.with_columns(max_col(["a", "b"]).alias("max_col"))\n

\n

输出。

\n

shape: (3, 3)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 a   \xe2\x94\x86 b    \xe2\x94\x86 max_col \xe2\x94\x82\n\xe2\x94\x82 --- \xe2\x94\x86 ---  \xe2\x94\x86 ---     \xe2\x94\x82\n\xe2\x94\x82 i64 \xe2\x94\x86 i64  \xe2\x94\x86 str     \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 1   \xe2\x94\x86 4    \xe2\x94\x86 b       \xe2\x94\x82\n\xe2\x94\x82 8   \xe2\x94\x86 5    \xe2\x94\x86 a       \xe2\x94\x82\n\xe2\x94\x82 3   \xe2\x94\x86 null \xe2\x94\x86 a       \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n