如何根据列的相等部分将切片分配给pandas数据框？

Question

I have sorted a roughly 1 million row dataframe by a certain column. I would like to assign groups to each observation based on equal sums of another column but I'm not sure how to do this.

Example below:

import pandas as pd
value1 = [25,27,20,22,28,20]
value2 = [.34,.43,.54,.43,.5,.7]

df = pd.DataFrame({'value1':value1,'value2':value2})

df.sort_values('value1', ascending = False)

df['wanted_result'] = [1,1,1,2,2,2]

Like this example, I want to sum my column (example column value1 ) and assign groups to have as close to equal value1 sums as they can. Is there a build in function to this?

Answer 1

Greedy Loop

Using Numba's JIT to quicken it up.

from numba import njit

@njit
def partition(c, n):
    delta = c[-1] / n
    group = 1
    indices = [group]
    total = delta

    for left, right in zip(c, c[1:]):
        left_diff = total - left
        right_diff = total - right
        if right > total and abs(total - right) > abs(total - left):
            group += 1
            total += delta
        indices.append(group)

    return indices

df.assign(result=partition(df.value1.to_numpy().cumsum(), n=2))

   value1  value2  result
4      28    0.50       1
1      27    0.43       1
0      25    0.34       1
3      22    0.43       2
2      20    0.54       2
5      20    0.70       2

---

This is NOT optimal. This is a greedy heuristic. It goes through the list and finds where we step over to the next group. At that point it decides whether it's better to include the current point in the current group or the next group.

This should behave pretty well except in cases with huge disparity in values with the larger values coming towards the end. This is because this algorithm is greedy and only looks at what it knows at the moment and not everything at once.

But like I said, it should be good enough.

Answer 2

I think, this is a kind of optimalisation problem (non-linear) and Pandas is definitively not any good candidate to solve it.

The basic idea to solve the problem can be as follows:

1. Definitions:

   • n - number of elements,
   • groupNo - the number of groups to divide into.

2. Start from generating an initial solution , eg take consecutive groups of n / groupNo elements into each bin .

3. Define the goal function , eg sum of squares of differences between sum of each group and sum of all elements / groupNo .

4. Perform an iteration:

   • for each pair of elements a and b from different bins, calculate the new goal function value, if these elements were moved to the other bin,
   • select the pair that gives the greater improvement of the goal function and perform the move (move a from its present bin to the bin, where b is, and vice versa).

5. If no such pair can be found, then we have the final result.

Maybe someone will propose a better solution, but at least this solution is some concept to start with.