熊猫数据框中一列的总和值按另一列排序，最高值为 N

Question

假设我有一个这样的列（基于对 B 列应用的某种排序）：

   A   B
0  2   5   
1  5   4
2  9   3
3  4   3
4  3   1

现在我必须找出不超过 15 的总和，但数据框中行的顺序不能改变（它需要按 B 列值排序）（虽然在计算总和时可以省略一行），所以在这种情况下row sum(0, 1 , 3 , 4)over column A = 14. 基本上是对一列的条件求和。

预期输出：df，其列 A 添加行

注意：

我只能跳过值 A > total_left 的那些行，必须始终仅按该顺序考虑这些行（由 B 优先），例如，数据帧中可能有数百万行，我必须绝对考虑值为

   A  B
0  5  5
1  5  4
3  5  3
4  1  1

我不能先取第 4 行，因为它的价值较小，我需要取 row(0)、row(1)、row(2)，因此无法对 column(A) 进行排序。

Answer 1

IIUC：

df = df[df['A'].sort_values().cumsum() < 15]

OUTPUT

   A  B
0  2  5
1  5  4
3  4  3
4  3  1

Answer 2

新答案

OP 遗漏了一个关键细节：如果它们对累积总和的贡献使其不超过最大值，则不能跳过按顺序排列的行。这大大改变了问题，无法再使用原始答案。

事实上，在当前的pandas 或numpy 中没有矢量化操作可以在预先未知的位置处实现这样的累积和。

如in this SO answer 所示，留给大型数组的最佳方法是使用numba。

这里是如何做到这一点，适应这个问题：

from numba import njit
from numba.types import bool_

@njit
def conditional_cumsum(x, skip_if_higher_than):
    total = 0
    result = np.zeros_like(x, dtype=bool_)
    for i, y in enumerate(x):
        if total + y <= skip_if_higher_than:
            total += y
            result[i] = True
    return result

示例

df = pd.DataFrame({'A': [2, 5, 9, 4, 3], 'B': [5, 4, 3, 3, 1]})
mask = conditional_cumsum(df['A'].values, skip_if_higher_than=15)

>>> df.loc[mask]
   A  B
0  2  5
1  5  4
3  4  3
4  3  1


df = pd.DataFrame({'A': [5, 5, 5, 1], 'B': [5, 4, 3, 1]})

>>> conditional_cumsum(df['A'].values, skip_if_higher_than=15)
array([ True,  True,  True, False])

速度

n = 1_000_000
np.random.seed(0)
df = pd.DataFrame({
    'A': np.random.uniform(size=n),
    'B': np.arange(n)
})

%timeit conditional_cumsum(df['A'].values, skip_if_higher_than=1)
460 µs ± 751 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit conditional_cumsum(df['A'].values, skip_if_higher_than=100)
460 µs ± 939 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

>>> df.loc[conditional_cumsum(df['A'].values, skip_if_higher_than=1)]
                   A       B
0       5.488135e-01       0
4       4.236548e-01       4
16      2.021840e-02      16
99      4.695476e-03      99
757     1.383350e-03     757
821     5.459649e-04     821
1070    6.642186e-04    1070
84341   3.310554e-06   84341
131245  1.989694e-05  131245
661553  7.071203e-07  661553

请注意：任何njit 函数都有一个初始“编译时间”。首先在一个小数组上运行它以让它发生，然后在大数组上运行。

原答案

由于您可以跳过行，因此选择列的顺序并不重要（我们可以在之后恢复它）。正如@MuhammadHassan 所说，最大子集将是A 的部分，已排序，总和为15：

s = df['A'].sort_values().cumsum() <= 15
idx = df.index.intersection(s[s].index)

>>> idx.tolist()
[0, 1, 3, 4]

# and
>>> df.loc[idx]
   A  B
0  2  5
1  5  4
3  4  3
4  3  1

原始答案注释

出于教学目的，我将保留此原始答案，但@MuhammadHassan 的答案是正确且更简洁的。为了防止UserWarning: Boolean Series key will be reindexed to match DataFrame index（并选择最多 15，这意味着最多并包括15）：

>>> df.loc[df['A'].sort_values().cumsum() <= 15]
   A  B
0  2  5
1  5  4
3  4  3
4  3  1