从两个数组中选择N个数字，其中k个具有最高总值

Question

假设 A、B 和 C 是三个数组，每个数组都包含 N 个数字：

 A = a[0], a[1], a[2], ..., a[N-1] 
 B = b[0], b[1], b[2], ..., b[N-1]
 C = c[0], c[1], c[3], ..., c[N-1] 

我想从数组 A 和数组 B 中选择最好的 k= a[i]+b[i] ，而不是 a[i]+b[i]。

乍一看，这似乎很简单，但我考虑得越多，它就变得越复杂。

最终我要在 Python 中实现此算法，但现阶段我只是试图了解如何编写高效的算法。

---

示例

为了澄清问题，算法的输入包括三个 N x 1 数组 A、B 和 C，以及一个整数值 k。其期望的输出是两个 k x 1 的索引列表，定义了来自 A、B（和C）的价值最大化的元素组合。

例如，假设 k=2，N=4, 则

 A = a[0], a[1], a[2], a[3] = 3, 1, 1, 0   
 B = b[0], b[1], b[2], b[3] = 1, 3, 0, 1  
 C = c[0], c[1], c[2], c[3] = 4, 4, 3, 2

即使在这个简单的例子中，有许多可能的组合。例如，如果从A选择元素i = 0,2，并从B选择元素j = 1,3，那么总值将是a [0] + a [2] + b [1] + b [3] = 8 。

另一方面，如果从A和B中选择元素i = 0,1和j = 0,1，则应用特殊扭曲：总值不是a [0] + a [1] + b [0] + b [1]，而是由c [0] + c [1] = 8给出。

在这个例子中，最大化总值的元素组合由A中的i = 0,2元素和B中的j = 1,2元素给出。这产生了一个总价值为a [0] + b [1] + c [2] = 9，可以验证它比任何其他组合都更大。

答案比较

以下是对三个提交的解决方案的快速比较。首先，我检查了它们所有的结果，它们都给出了预期的结果。顺便说一句，它们都不要求C的元素弱大于A和B中相应元素的总和，因此我在我的性能评估中取消了这个假设。

import numpy as np
from utils import tic, toc  # simple wrapper to time.perf_counter()

k, N = 10, 1000

A = list(np.random.random_sample([N]))
B = list(np.random.random_sample([N]))
C = list(np.random.random_sample([N]))

tic()
print(optimal_choices(k, A, B, C))  # solution by btilly
toc()

tic()
print(maxPicks(A.copy(), B.copy(), C.copy(), k))  # solution by Eric T-M
toc()

tic()
print(maxSum(A, B, C, k))  # solution by Alain T.
toc()

我尝试了各种不同的k和N的组合，似乎@btilly的算法只要k很小，就能很好地处理大规模数据N。相反地，当k相对于N很大时，@Alain-T.的算法表现得更好。总体而言，@Eric-T-M的算法表现最好，在k和N的扩展性方面都表现良好。

小规模问题: k = 10 and N = 500

• btilly's算法: 0.49秒
• Eric T-M's算法: 0.00秒
• Alain T.'s算法: 0.52秒

小k，大N: k = 10 and N = 1000

• btilly's算法: 0.89秒
• Eric T-M's算法: 0.00秒
• Alain T.'s算法: 1.99秒

大k，小N: k = 80 and N = 100

• btilly's算法: 1.54秒
• Eric T-M's算法: 0.00秒
• Alain T.'s算法: 0.09秒

中等规模问题: k = 50 and N = 1000

• btilly's算法: 13.01秒
• Eric T-M's算法: 0.00秒
• Alain T.'s算法: 8.55秒

大规模问题1: k = 10 and N = 1_000_000

• Eric T-M's算法: 1.03秒

大规模问题2: k = 1_000 and N = 100_000

• Eric T-M's算法: 10.22秒

（为了进行基准测试，我删除了Alain T.代码中的排序操作，使之具有可比性。）

- Bob

---

元素 c 是否始终大于元素 a 和 b 对应元素之和？ - gionni

那么，在这种情况下，你应该总是从C[]中进行选择。 - wildplasser

@gionni 是的，在我的应用程序中，我可以假设无失一般性地认为 c 中的元素弱大于其对应的 a 和 b 元素之和（抱歉来回反复，我在@wildplasser的下面评论中稍微有点困惑） - Bob

@wildplasser：不一定，可以看看我上面的例子。 - Bob

Answer 1

尝试这个。它需要 O(N^2) 时间，而且相当简单。

def maxPicks(A,B,C,k):
    # returns the tuple (list of entries picked in A, list of entries picked in B, total value)

    # BASE CASE
    if k == 0:
        return ([], [], 0)
    aMax = max(A)
    bMax = max(B)
    cMax = max(C)

    if (aMax + bMax) > cMax:
        aIdx = A.index(aMax)
        bIdx = B.index(bMax)
        B[aIdx] = C[aIdx] - A[aIdx]
        A[aIdx] = -2
        C[aIdx] = -1
        A[bIdx] = C[bIdx] - B[bIdx]
        B[bIdx] = -2
        C[bIdx] = -1
        nextPicks = maxPicks(A,B,C,k-1)
        return (nextPicks[0] + [aIdx], nextPicks[1] + [bIdx], nextPicks[2] + aMax + bMax)
    else:
        cIdx = C.index(cMax)
        A[cIdx] = -1
        B[cIdx] = -1
        C[cIdx] = -1
        nextPicks = maxPicks(A,B,C,k-1)
        return (nextPicks[0] + [cIdx], nextPicks[1] + [cIdx], nextPicks[2] + cMax)

这是它的工作方式：

基本情况应该是不言自明的。否则，我们将比较在A中所有条目的最大值和B中所有条目的最大值之和与C中所有条目的最大值。如果这个和较大，那么安全起见，我们可以从A和B中选择这些条目，但在进行更多选择之前，我们需要将我们选择的条目及其对应的C中的条目设置为负值。顺便说一下，我假设A、B和C中的所有值都是非负数，因此通过将它们置为负数，我们禁止算法再次选择它们。如果这个假设是错误的，你可能想将这些值设置为极低的负数以禁止重复选择。我们还看到，如果我们选择了A[i]，则B[i]的值现在是C[i]-A[i]，因为选择B[i]会使我们失去A[i]中的值并获得C[i]中的值，对于条目A[j]也是同样道理，如果我们选择了B[j]。

另一方面，如果C中的最大条目大于等于aMax+bMax，我们想选择它（通过选择A和B中的相应条目），因为没有其他的A和B条目或仅C条目的选择会更有价值。此时，我们知道我们不想再次选择A[i]、B[i]或C[i]，所以我们将它们全部置为负值。

- Eric T-M

Answer 2

这可以用动态规划来解决。

# Helper function to get out of the data structure.
def get_nested_array (data, path):
    for x in path:
        if data is None or len(data) <= x:
            return None
        else:
            data = data[x]
    return data

# Helper function to set data in the data structure.
def set_nested_array (data, path, value):
    # Navigate there
    for x in path[0:len(path)-1]:
        while len(data) <= x:
            data.append([])
        if data[x] is None:
            data[x] = []
        data = data[x]
    while len(data) <= path[-1]:
        data.append(None)
    data[path[-1]] = value

# Is this option better than what is there?  If so, then add it.
def possibly_add_choice (best_choice, pos, i, j, current_sum, last_i, last_j):
    best_prev = get_nested_array(best_choice, [pos, i, j])
    if best_prev is None or best_prev[0] < current_sum:
        set_nested_array(best_choice, [pos, i, j], (current_sum, last_i, last_j))

# Our program.
def optimal_choices (k, A, B, C):
    # best_choice[pos][i][j] = (max_sum, last_a, last_b)
    # where:
    #    We have made all choices in the range 0..pos-1
    #    We chose i out of A
    #    We chose j out of B
    # and
    #    max_sum is the best possible sum
    #    last_a is the last index chosen from a
    #    last_b is the last index chosen from b
    # then we can find the answer working backwards from
    # best_choice[len(A)][k][k]
    #
    best_choice = []

    # Enter the empty set answer
    set_nested_array(best_choice, [0, 0, 0], (0, None, None))

    for pos in range(len(A)):
        best_choice.append([])
        best_choice_for_pos = best_choice[pos]
        for i in range(k+1):
            if len(best_choice_for_pos) <= i:
                break
            best_choice_for_i = best_choice_for_pos[i]
            for j in range(k+1):
                if len(best_choice_for_i) <= j:
                    break
                last_sum, last_i, last_j = best_choice_for_i[j]

                # Try all 4 things we can choose here.  Nothing, or A or B or both.
                possibly_add_choice(best_choice, pos+1, i, j, last_sum, last_i, last_j)
                possibly_add_choice(best_choice, pos+1, i+1, j, last_sum + A[pos], pos, last_j)
                possibly_add_choice(best_choice, pos+1, i, j+1, last_sum + B[pos], last_i, pos)
                possibly_add_choice(best_choice, pos+1, i+1, j+1, last_sum + C[pos], pos, pos)

    # Now we have the answer, it is just a question of decoding it.
    if get_nested_array(best_choice, [len(A), k, k]) is None:
        return (None, None)
    else:
        choose_a = []
        choose_b = []
        best_spot = [len(A), k, k]
        max_sum, last_i, last_j = get_nested_array(best_choice, best_spot)
        while last_i is not None or last_j is not None:
            # Figure out where we last had a choice and what was chosen.
            if last_i is None:
                last_pos = last_j
                i_dec = 0
                j_dec = 1
            elif last_j is None:
                last_pos = last_i
                i_dec = 1
                j_dec = 0
            else:
                last_pos = max(last_i, last_j)
                i_dec = 0
                j_dec = 0
                if last_pos == last_i:
                    i_dec = 1
                if last_pos == last_j:
                    j_dec = 1

            # record the choice.
            if 1 == i_dec:
                choose_a.append(last_pos)
            if 1 == j_dec:
                choose_b.append(last_pos)

            # Go back to that spot
            max_sum, last_i, last_j = get_nested_array(best_choice, [last_pos, k-len(choose_a), k-len(choose_b)])

        # We walked backwards to generate these lists.  So they are currently reversed.
        return (list(reversed(choose_a)), list(reversed(choose_b)))

print(optimal_choices(2, [3, 1, 1, 0 ], [1, 3, 0, 1], [4, 4, 3, 2]))

- btilly

Answer 3

如果您将A和B展开为索引对的列表，并给出它们各自的总和（应用C的例外情况），则可以迭代地取最大总和并在每个步骤中排除相应的对。这将选择每次迭代中剩余对中可能的最高总和：

def maxSum(A,B,C,K):
    S = [ ([a+b,C[i]][i==j],i,j) for i,a in enumerate(A)  
                                 for j,b in enumerate(B)]
    usedA,usedB  = set(),set()
    for _ in range(K):
        _,i,j = max(s for s in S if not(s[1] in usedA or s[2] in usedB))
        usedA.add(i)
        usedB.add(j)
    return sorted(usedA),sorted(usedB)

输出：

A = [3, 1, 1, 0]   
B = [1, 3, 0, 1]  
C = [4, 4, 3, 2]
print(maxSum(A,B,C,2)) # ([0, 2], [1, 2])   total is 9

A = [1,2,3,4]
B = [4,5,6,2]
C = [5,9,9,6]

print(maxSum(A,B,C,2)) # ([1, 3], [1, 2])  total is 19

print(maxSum(A,B,C,3)) # ([1, 2, 3], [0, 1, 2]) total is 26

- Alain T.

---

不起作用。尝试 A = [4, 1, 2, 0]，B = [4, 3, 0, 5]，C = [8, 4, 2, 8] 和 k = 2。你正在寻找 ([0, 2], [0, 3])，但忽略了 ([0, 3], [0, 3]) 更好的选择。将 C 更改为 [6, 4, 2, 8]，您会得到相同的答案，并且会错过选择两个 0 现在更糟的情况。 - btilly

你说得对，它首先找到了12，但由于前一个选择已经排除了第二个最佳总和的索引，所以无法再计算出8的总和（在C中）。要解决这个问题需要一些动态编程，但我现在没有时间实现，不过我相信配对方法仍然可能有用。 - Alain T.