python - 极坐标 DataFrame 行之间的分割值

Question

我想找到一种使用极坐标将 DataFrame 的值分布在另一个 DataFrame 的行之间的方法(无需迭代行)。

我有一个数据框，其中包含要分配的金额:

<表类=“s-表”> <标题> 姓名 金额 <正文> 一个 100 B 300 C 250

以及一个目标 DataFrame，我想使用通用的“名称”列将分布式值(在新列中)附加到该目标 DataFrame。

<表类=“s-表”> <标题> 姓名 项目 价格 <正文> 一个 x1 40 一个 x2 60 B y1 50 B y2 150 B y3 200 C z1 400

目标中的行已排序，分配的金额应与每行中的价格匹配(只要有足够的剩余金额)。

因此本例中的结果应如下所示:

<表类=“s-表”> <标题> 姓名 项目 价格 分配金额 <正文> 一个 x1 40 40 一个 x2 60 60 B y1 50 50 B y2 150 150 B y3 200 100 C z1 400 250

在此示例中，我们可以分配 A 的金额，使它们与价格相同。然而，对于 B 和 C 的最后一项，我们写下剩余的金额，因为价格太高。

有没有有效的方法来做到这一点？

我最初的解决方案是计算目标数据帧中新列中价格的累积和，然后左连接源数据帧并减去累积和的值。如果金额足够高，这将起作用，但对于 B 和 C 的最后一项，我会得到负值，而不是剩余金额。

编辑

示例数据框:

import polars as pl

df1 = pl.DataFrame({"Name": ["A", "B", "C"], "Amount": [100, 300, 250]})
df2 = pl.DataFrame({"Name": ["A", "A", "B", "B", "B", "C"], "Item": ["x1", "x2", "y1", "y2", "y3", "z"],"Price": [40, 60, 50, 150, 200, 400]})

Answer 1

@jquurious，很好的答案。这可能稍微简洁一些:

(
    df2.join(df1, on="Name")
    .with_columns(
        pl.min([
            pl.col('Price'),
            pl.col('Amount') -
            pl.col('Price').cumsum().shift_and_fill(1, 0).over('Name')
        ])
        .clip_min(0)
        .alias('assigned')
    )
)

shape: (6, 5)
┌──────┬──────┬───────┬────────┬──────────┐
│ Name ┆ Item ┆ Price ┆ Amount ┆ assigned │
│ ---  ┆ ---  ┆ ---   ┆ ---    ┆ ---      │
│ str  ┆ str  ┆ i64   ┆ i64    ┆ i64      │
╞══════╪══════╪═══════╪════════╪══════════╡
│ A    ┆ x1   ┆ 40    ┆ 100    ┆ 40       │
│ A    ┆ x2   ┆ 60    ┆ 100    ┆ 60       │
│ B    ┆ y1   ┆ 50    ┆ 300    ┆ 50       │
│ B    ┆ y2   ┆ 150   ┆ 300    ┆ 150      │
│ B    ┆ y3   ┆ 200   ┆ 300    ┆ 100      │
│ C    ┆ z    ┆ 400   ┆ 250    ┆ 250      │
└──────┴──────┴───────┴────────┴──────────┘