Início / coding / Perguntas / 79576094
Accepted
crazydragon777
Asked: 2025-04-16 06:09:30 +0800 CST

Qual é a melhor maneira de agrupar por mais de uma coluna?

  • 772

Basicamente, quero agrupar todos os registros que correspondem a qualquer uma das chaves

Entrada:

|id|key1|key2|
|-|-|-|
|1|a|x|
|2|a|y|
|3|b|y|
|4|c|z|

Saída desejada: (As 3 primeiras linhas são agrupadas)

|key1 (or any identifier, I only care about the final counts)|len|
|-|-|
|a|3|
|c|1|

Minha tentativa de solução (incompleta):

import polars as pl

df = pl.DataFrame({
    'id': [1, 2, 3, 4],
    'key1': ['a', 'a', 'b', 'c'],
    'key2': ['x', 'y', 'y', 'z'],
})
df = (
    df.group_by(
        'key1',
        'key2',
    )
    .len()
    .group_by('key1')
    .agg(
        pl.col('key2'),
        pl.col('len').sum()
    )
)
print(df)

O que dá:

|key1|key2|len|
|-|-|-|
|a|["y", "x"]|2|
|b|["y"]|1|
|c|["z"]|1|

No entanto, não tenho certeza de como agruparia isso ainda mais por key2 (mesclando as linhas a e b, já que elas têm um valor comum de y) preservando a soma de len

python

1 respostas

  1. Best Answer

    Se entendi corretamente, você quer criar relações de equivalência entre key1e key2e continuar mesclando coisas em grupos se houver uma cadeia de relacionamentos que conecte os dois? Então, neste caso, como há uma linha, a, yquaisquer linhas que contenham aou ydevem ser mescladas?

    Neste caso é difícil expressar isso diretamente em Polars, mas você pode fazer um híbrido com scipy's DisjointSet.

    Primeiro, reúna todas as chaves e relacionamentos exclusivos:

    keys = df.select(pl.col.key1.append(pl.col.key2).unique())
edges = df.select(["key1", "key2"]).unique()

    Em seguida, insira-os em um DisjointSete mescle-os em grupos:

    from scipy.cluster.hierarchy import DisjointSet

ds = DisjointSet(keys.to_series().to_list())
for a, b in edges.iter_rows():
    ds.merge(a, b)

    Em seguida, converta isso DisjointSetem uma DataFramechave de mapeamento para agrupar o índice:

    group_idx = (
    pl.DataFrame({"key": [list(s) for s in ds.subsets()]})
        .lazy()
        .with_row_index("group_idx")
        .explode("key")
)

    Por fim, juntamos o group_idxao nosso dataframe original e agrupamos pelo índice do grupo:

    (df.lazy()
    .join(group_idx, left_on="key1", right_on="key")
    .group_by("group_idx")
    .len()
    .collect())

    Dando assim a nossa saída:

    ┌───────────┬─────┐
│ group_idx ┆ len │
│ ---       ┆ --- │
│ u32       ┆ u32 │
╞═══════════╪═════╡
│ 0         ┆ 1   │
│ 1         ┆ 3   │
└───────────┴─────┘
    • 2

relate perguntas