给定二维数组，长度为 K 且总和为 S 的子序列数

Question

我希望在给定一个数组的情况下，找出总和为 S 的长度为 K 的子序列的数量。

示例输入：

a=[1,1,1,2,2] & K=2 & S=2

样本输出：

3 {because a[0],a[1]; a[1]a[2]; a[0]a[2] are only three possible for the case}

我尝试在 Python 中为初学者编写一个递归循环，但它没有按预期提供输出。请你帮我找到我可能存在的漏洞错过了。

def rec(k, sum1, arr, i=0):                                 
    #print('k: '+str(k)+' '+'sum1: '+str(sum1))     #(1) BaseCase: 
    if(sum1==0 and k!=0):                           #    Both sum(sum1) required and 
        return 0                                    #    numbers from which sum is required(k)
    if(k==0 and sum1 !=0):                          #    should be simultaneously zero
        return 0                                    #    Then required subsequences are 1
    if(k==0 and sum1==0 ):                          # 
        return 1                                    #
    base_check = sum1!=0 or k!=0                    #(2) if iterator i reaches final element 
    if(i==len(arr) and base_check):                 #    in array we should return 0 if both k 
        return 0                                    #    and sum1 aren't zero
                                                    #    func rec for getting sum1 from k elements 
    if(sum1<arr[0]):                                #    takes either first element or rejects it
        ans=rec(k-1,sum1,arr[i+1:len(arr)],i+1)     #    so 2 cases in else loop 
        print(ans)                                  #    i is taken in as iterator to provide array 
    else:                                           #    input to rec func from 2nd element of array 
        ans=rec(k-1, sum1-arr[0], arr[i+1:len(arr)],i+1)+rec(k, sum1, arr[i+1:len(arr)],i+1)
        #print('i: '+str(i)+' ans: '+str(ans))      
    return(ans)

a=[1,1,1,2,2]
print(rec(2,2,a))

我仍然无法处理如何进行更改。一旦编写了这个正常的递归代码，我可能会采用 DP 方法。

Answer 1

使用itertools.combinations

函数itertools.combinations 返回给定长度的所有子序列。然后我们过滤以仅保留总和为所需值的子序列。

import itertools
def countsubsum(a, k, s):
  return sum(1 for c in itertools.combinations(a,k) if sum(c)==s)

修复代码

您的代码看起来不错，但有两点看起来有问题。

这是什么if？

起初我对if(sum1<arr[0]): 有点困惑。我认为你可以（并且应该）总是去 else 分支。在考虑了更多之后，我知道如果arr[0] 太大而无法采用，您正试图摆脱两个递归调用之一，这很聪明，但这假设数组中的所有元素都是非负的.如果允许数组包含负数，那么你可以在子序列中包含一个大的a[0]，并希望有一个负元素来弥补。所以如果数组可以包含负数，你应该去掉这个if/else，并始终从else分支执行两个递归调用。

你切错了

您维护一个变量i 来记住数组中的起始位置；但你也对数组进行切片。很快你的索引就出错了。您应该使用切片，或者使用索引i，但不能同时使用。

# WRONG
ans=rec(k-1, sum1-arr[0], arr[i+1:len(arr)],i+1)+rec(k, sum1, arr[i+1:len(arr)],i+1)

# CORRECT
ans = rec(k-1, sum1-arr[i], arr, i+1) + rec(k, sum1, arr, i+1)

# CORRECT
ans = rec(k-1, sum1-arr[0], arr[1:]) + rec(k, sum1, arr[1:])

要了解为什么同时使用切片和索引会产生错误结果，请运行以下代码：

def iter_array_wrong(a, i=0):
  if (a):
    print(i, a)
    iter_array_wrong(a[i:], i+1)

def iter_array_index(a, i=0):
  if i < len(a):
    print(i, a)
    iter_array_index(a, i+1)

def iter_array_slice(a):
  if a:
    print(a)
    iter_array_slice(a[1:])

print('WRONG')
iter_array_wrong(list(range(10)))
print()
print('INDEX')
iter_array_index(list(range(10)))
print()
print('SLICE')
iter_array_slice(list(range(10)))

另请注意，a[i:len(a)] 完全等同于 a[i:]，a[0:j] 等同于 a[:j]。

简洁版的递归

递归计算使用数组第一个元素的子序列，以及不使用数组第一个元素的子序列，并将两个计数相加。为了避免重复显式地对数组进行切片，这是一项昂贵的操作，我们保留了一个变量 start 以记住我们只处理子数组 a[start:]。

def countsubsum(a, k, s, start=0):
  if k == 0:
    return (1 if s == 0 else 0)
  elif start == len(a):
    return 0
  else:
    using_first_element = countsubsum(a, k-1, s-a[start], start+1)
    notusing_first_elem = countsubsum(a, k, s, start+1)
    return using_first_element + notusing_first_elem