按另一列排序的pandas数据帧中一列的和值，直到值N

Question

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

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

现在，我必须找出15的总和，但数据帧中行的顺序不能改变(它需要按列B的值排序((在计算总和时可以省略一行(，所以在这种情况下，列a上的行总和(0，1，3，4(=14。基本上是一列上的条件求和。

预期输出：df，其列A添加了行<=N(此处为15(。

注：

我只能跳过那些值A>total_left，必须始终只按该顺序考虑这些行(按B排序(，例如，数据帧中可能有数百万行，我必须肯定地考虑其值<total_left(total-sum_till_now(，我不能离开这些行，例如，我有一个类似的东西

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

我不能先取第4行，因为它的值较小，我需要取第(0(行、第(1(行和第(2(行，因此无法在列(A(上排序。

Answer 1

新答案

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

事实上，在当前pandas或numpy中不存在矢量化操作，该矢量化操作将在预先未知的位置重置的情况下实现这样的累积和。

正如这个SO答案中所指出的，留给大型阵列的最佳方法是使用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

请注意：有一个首字母"编译时"；对于任何CCD_ 4函数。先在一个小数组上运行它，然后在大数组上运行。

原始答案

由于您可以跳过行，所以选择列的顺序并不重要(我们可以稍后恢复(。正如@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

Answer 2

IIUC:

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

OUTPUT

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