列表选k个元素求最大乘积的Python代码Bug排查

题目信息 #

本次调试对应竞技编程题目为Codeforces平台第1696场竞赛的H题。

现有实现代码 #

def max_prod(ar,k):

    a = sorted(ar, key=abs)[::-1]
    n = len(ar)
    ## assuming n>k

    ans = 1
    neg = 0
    pos = 0
    last_neg = None
    last_pos = None
    for i in range(k):
        if a[i] <0:
            neg+=1
            last_neg = i
        elif a[i] >0:
            pos+=1
            last_pos =i
        ans *= a[i]

    if ans>=0:
        return ans
    else:

        # option A - remove last negative number and replace with next positive number

        # option B - remove last postive number and replace wit next negative number
        next_non_neg_found = False
        next_neg_found = False
        next_non_neg_ind = None
        next_neg_ind = None
        next_zero_found = False
        next_zero_ind = None

        for j in range(k,n):

            if a[j] >= 0 and next_non_neg_found==False:
                next_non_neg_ind = j
                next_non_neg_found = True

            if a[j] < 0 and next_neg_found==False:
                next_neg_ind = j
                next_neg_found = True

        op1,op2 = None,None

        #Implement option A
        if next_non_neg_found:
            if last_neg is not None:    ## last_neg cannot be none because ans =-ve
                op1 = (ans*a[next_non_neg_ind])//a[last_neg]
        #else: ##if no postive next then cannot do anything

        #Implemnet option B
        if next_neg_found:
            if last_pos is not None:
                op2 = (ans*a[next_neg_ind])//a[last_pos]
            #else - nothing can be done
        #else: - nothing

        if op1 is not None and op2 is not None:
            return max(ans,op1,op2)
        else:
            if op1 is not None:
                return max(ans,op1)
            elif op2 is not None:
                return max(ans,op2)
            else:
                return ans

def powerset(seq):
    """
    Returns all the subsets of this set. This is a generator.
    """
    if len(seq) <= 1:
        yield seq
        yield []
    else:
        for item in powerset(seq[1:]):
            yield [seq[0]]+item
            yield item


def calc_f(a,k,memo):
    n = len(a)

    if tuple(a) in memo:
        return memo[tuple(a)]
    else:

        if n<k:
            memo[tuple(a)] = 0
            return memo[tuple(a)]
        else:
            if n == k:
                prod = 1
                for i in range(0,n):
                    prod *= a[i]
                memo[tuple(a)] = prod
                return memo[tuple(a)]
            elif n>k:
                memo[tuple(a)] = max_prod(a,k)
                return memo[tuple(a)]


n,k = list(map(int,input().split()))

a = list(map(int,input().split()))

sub = [r for r in powerset(a)]

#print(len(sub))

sum = 0
memo = {}
for s in sub:
    sum += calc_f(s,k,memo)

print(sum%(10**9 +7))
print(memo)

问题描述 #

当前代码可以通过大部分测试用例，但在少数测试用例上返回错误结果。初步判断问题大概率出在max_prod函数中，存在遗漏的边界场景，排查一整天仍未定位问题。需要帮忙指出代码中的错误，同时确认当前采用的算法思路是否可行，若思路存在问题请予以纠正。

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