algorithm - How to divide a set into two sets such that the difference of the average is minimum?

Question

As I understand, it is related to the partition problem.

But I would like to ask a slightly different problem which I don't care about the sum but the average. In this case, it needs to optimize 2 constraints (sum and number of items) at the same time. It seems to be a harder problem and I cannot see any solutions online.

Are there any solutions for this variant? Or how does it relate to the partition problem?

---

Example:

input X = [1,1,1,1,1,6]
output based on sum: A = [1,1,1,1,1], B=[6]
output based on average: A = [1], B=[1,1,1,1,6]

question from:https://stackoverflow.com/questions/65932236/how-to-divide-a-set-into-two-sets-such-that-the-difference-of-the-average-is-min

Answer 1

On some inputs, a modification of the dynamic program for the usual partition problem will give a speedup. We have to classify each partial solution by its count and sum instead of just sum, which slows things down a bit. Python 3 below (note that the use of dictionaries implicitly collapses functionally identical partial solutions):

def children(ab, x):
    a, b = ab
    yield a + [x], b
    yield a, b + [x]


def proper(ab):
    a, b = ab
    return a and b


def avg(lst):
    return sum(lst) / len(lst)


def abs_diff_avg(ab):
    a, b = ab
    return abs(avg(a) - avg(b))


def min_abs_diff_avg(lst):
    solutions = {(0, 0): ([], [])}
    for x in lst:
        solutions = {
            (sum(a), len(a)): (a, b)
            for ab in solutions.values()
            for (a, b) in children(ab, x)
        }
    return min(filter(proper, solutions.values()), key=abs_diff_avg)


print(min_abs_diff_avg([1, 1, 1, 1, 1, 6]))