满足两组兼容性要求的m个元素最小成本算法问题排查

问题描述 #

给定长度n≥1的cost、compatible1、compatible2数组：

• cost[i]为元素i的成本；
• compatible1[i]/compatible2[i]为1表示元素i兼容对应组，0则不兼容。
  给定min_compatible，表示每组需满足的最低兼容元素数量，兼容两组的元素可同时计入两组配额。需返回满足兼容性要求的最小成本。

思路与实现 #

将元素分为仅兼容组1、仅兼容组2、兼容两组三类，用堆维护各类元素的最小成本。优先选择兼容两组的元素（当其成本低于选1个仅组1加1个仅组2元素的成本时），否则选当前最低成本的单组兼容元素，最后补充未满足的配额。实现代码如下：

from heapq import heappush, heappop

def getMinCost(cost, compatible1, compatible2, min_compatible):
    # set of indices compatible with each
    c1 = {i for i,c in enumerate(compatible1) if c == 1}
    c2 = {i for i,c in enumerate(compatible2) if c == 1}

    # not enough to fulfill reqs
    if len(c1) < min_compatible or len(c2) < min_compatible:
        return -1

    # make disjoint -> one, the other, or both
    c1, c2, c3 = c1 - c2, c2 - c1, c2 & c1

    # fill the disjoint heaps
    heap1 = []
    heap2 = []
    heap3 = []
    for i,c in enumerate(cost):
        if i in c3:
            heappush(heap3, c)
        elif i in c2:
            heappush(heap2, c)
        elif i in c1:
            heappush(heap1, c)

    # handle edge case (one empty early)
    heappush(heap1, float('inf'))
    heappush(heap2, float('inf'))
    heappush(heap3, float('inf'))

    # amount fulfilling each
    a1 = 0
    a2 = 0
    # minimum cost counter
    minCost = 0

    # both are below threshold, prioritize intersection
    while a1 < min_compatible and a2 < min_compatible:

        # if first cond true, advantageous to take the one
        # that fulfills both
        if heap1[0] + heap2[0] >= heap3[0]:
            minCost += heappop(heap3)
            a1 += 1
            a2 += 1
        elif heap1[0] <= heap2[0]:
            minCost += heappop(heap1)
            a1 += 1
        else:
            minCost += heappop(heap2)
            a2 += 1


    # deal with leftovers

    while a1 < min_compatible:

        if heap1[0] <= heap3[0]:
            minCost += heappop(heap1)
        else:
            minCost += heappop(heap3)
            a2 += 1

        a1 += 1

    while a2 < min_compatible:

        if heap2[0] <= heap3[0]:
            minCost += heappop(heap2)
        else:
            minCost += heappop(heap3)
            a1 += 1

        a2 += 1


    return minCost

问题 #

该代码仅通过15个测试用例中的8个，所有错误均为输出结果不正确，无法定位错误原因，需排查问题并避免后续重复犯错。

内容的提问来源于stack exchange，提问作者user129393192