首页 > 解决方案 > 将数量分配给数组元素以使对的差异最小的最快方法?

问题描述

给定一个数字数组arr和一个整数x,分配x使得任何对之间的差异尽可能小。

例如arr = [4,2,0] 和x = 10;

答案应该是 [6,5,5];

必须使用所有x

标签: algorithm

解决方案


将最终均值计算为(sum(arr) + x) / len(arr)如果我们也可以减少,那将是所有数字的理想目标。

四舍五入的商告诉我们每个数字应该变成的最小值,余数告诉我们有多少数字应该加上一个额外的 1。在消除已经太大的数字之后这样做。

总时间 O(n log n)。

Python实现:

def distribute(arr, x):
    total = sum(arr) + x
    I = sorted(range(len(arr)), key=arr.__getitem__)
    while I:
        minimum, additional = divmod(total, len(I))
        if arr[I[-1]] <= minimum:
            break
        total -= arr[I.pop()]
    for i in sorted(I):
        arr[i] = minimum
        if additional > 0:
            arr[i] += 1
            additional -= 1

测试一些硬编码输入、较大的随机输入和详尽的小输入的结果:

433103 tests passed
0 tests failed

完整代码(在线试用!):

from random import choices
from itertools import product

def distribute(arr, x):
    total = sum(arr) + x
    I = sorted(range(len(arr)), key=arr.__getitem__)
    while I:
        minimum, additional = divmod(total, len(I))
        if arr[I[-1]] <= minimum:
            break
        total -= arr[I.pop()]
    for i in sorted(I):
        arr[i] = minimum
        if additional > 0:
            arr[i] += 1
            additional -= 1

def naive(arr, x):
    for _ in range(x):
        arr[arr.index(min(arr))] += 1

passed = failed = 0

def test(arr, x):
    expect = arr.copy()
    naive(expect, x)
    result = arr.copy()
    distribute(result, x)
    global passed, failed
    if result == expect:
        passed += 1
    else:
        failed += 1
        print('failed:')
        print(f'{arr =    }')
        print(f'{expect = }')
        print(f'{result = }')
        print()

# Tests from OP, me, and David
test([4, 2, 0], 10)
test([4, 2, 99, 0], 10)
test([20, 15, 10, 5, 0], 10)

# Random larger tests
for x in range(1000):
    arr = choices(range(100), k=100)
    test(arr, x)

# Exhaustive smaller tests
for n in range(5):
    for arr in product(range(10), repeat=n):
        arr = list(arr)
        for x in range(n * 10):
            test(arr, x)

print(f'{passed} tests passed')
print(f'{failed} tests failed')

推荐阅读