如何从polar Dataframe中选择具有最大和的列?



我有极地数据框架与一堆列我需要找到,比如说,和最大的那一列下面的代码片段汇总了所有列,但是我错过了选择具有最大值的列的最后一步。我是一个极地新手,如果这不是很明显:)。当然,我可以很容易地提取到numpy并找到最大值等,但我正试图尽可能多地接受polar的功能

test=pl.DataFrame(
{
"a": [0, 1, 3, 4],
"b": [0, 0, 0, 0],
"c": [1, 0, 1, 0],
}
)
max_col=test.select(pl.col(test.columns).sum())  
shape: (1, 3)      
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 8   ┆ 0   ┆ 2   │

我会这样做,作为一个熔体/过滤器。

test 
.select(pl.all().sum()) 
.melt() 
.filter(pl.col('value')==pl.col('value').max())

如果你想要原始形状,那么单链有点硬。我就这样做。

allcalc=test 
.select(pl.all().sum())
allcalc.select(allcalc.melt().filter(pl.col('value')==pl.col('value').max()) 
.get_column('variable').to_list())

如果有平局,上面的方法可以使用,例如:

test=pl.DataFrame(
{
"a": [0, 1, 3, 4],
"b": [0, 0, 0, 0],
"c": [1, 0, 1, 6],
}
)

那么在任何情况下你都会得到'a'和'c'。

还有一个选择是做一个sum和一个transposetranspose是一个昂贵的操作,但这里我们在已经聚合的数据帧上使用它,所以影响应该不会很大。

作为一个优点,它使代码更短:

test.select(
test.sum().transpose(include_header=True)
.sort('column_0', descending=True).head(1)[0,0]
)
# Result
shape: (4, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 0   │
│ 1   │
│ 3   │
│ 4   │
└─────┘

您可以使用polars.DataFrame.row作为字典(row(0, named=True))来获取列和。然后使用max查找和最高的列。

test = pl.DataFrame(
{
"a": [0, 1, 3, 4],
"b": [0, 0, 0, 0],
"c": [1, 0, 1, 0],
}
)
col_sum_dict = test.select(pl.all().sum()).row(0, named=True)
res = test.select(
[col for col, col_sum in col_sum_dict.items() 
if col_sum == max(col_sum_dict.values())]
)

输出:

>>> res
shape: (4, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 0   │
│ 1   │
│ 3   │
│ 4   │
└─────┘
>>> col_sum_dict
{'a': 8, 'b': 0, 'c': 2}

请注意,如果存在平局,即多个列的和最高,则此方法有效。例如:

test = pl.DataFrame(
{
"a": [0, 1, 3, 4],
"same_as_a": [0, 1, 3, 4],
"b": [0, 0, 0, 0],
"c": [1, 0, 1, 0],
}
)

在搜索结果

>>> res
shape: (4, 2)
┌─────┬───────────┐
│ a   ┆ same_as_a │
│ --- ┆ ---       │
│ i64 ┆ i64       │
╞═════╪═══════════╡
│ 0   ┆ 0         │
│ 1   ┆ 1         │
│ 3   ┆ 3         │
│ 4   ┆ 4         │
└─────┴───────────┘