leetcode上subset sum问题总结

本文是关于leetcode上subset sum类问题的总结，目前涉及到的题目有：

416 partition equal subset sum 问题
698 partition to k equal sum subset 问题

partition equal subset sum

Given a non-empty array containing only positive integers, find if the array can be partitioned into two
subsets such that the sum of elements in both subsets is equal.

Note:
Each of the array element will not exceed 100.
The array size will not exceed 200.
Example 1:

Input: [1, 5, 11, 5]

Output: true

Explanation: The array can be partitioned as [1, 5, 5] and [11].
Example 2:

Input: [1, 2, 3, 5]

Output: false

Explanation: The array cannot be partitioned into equal sum subsets.

题目的意思是，给定给一个数组，都是正数，判断是否能分成两个子数组，使得两个子数组的是一样的，换言之，能否在数组中找到一个或多个数，使得其和为原数组和的一半。这其中隐藏的两个条件就是原数组的和必须能被2整除以及最大的数不能超过原数组和的一半。我们可以用这两个条件来做一些边缘情况的处理。

我们很容易想到，如果能够分成这样的两个子数组，那么对于每个原数组中的数，非此即彼，我们可以用回溯法来试探性的探索每个元素的位置。也就是遍历每个元素，放到和还不够的子数组中，直到满足回溯的条件：遇到一个元素没有数组可以放。下面是这个思路的代码：

def canPartition(self, nums: 'List[int]') -> 'bool':
    target, remainder = divmod(sum(nums), 2)
    if remainder or max(nums) > target:
        return False
    
    if target in nums: return True

    def search(groups):
        if not nums: return True
        v = nums.pop()
        for i, g in enumerate(groups):
            if g + v <= subsum:
                groups[i] += v
                if search(groups): return True
                groups[i] -= v
        nums.append(v)
        return False 

    return search([0] * 2)

第二种方法是动态规划，一般能用回溯法也可以用动态规划。而动态规划的重点在于找到状态和状态转移公式。前面已经讲到，题意可以理解成找到若干个元素的和等于原数组和的一半。那么我们可以维护一个状态数组lookup，其中lookup[i]为真表示可以找到若干个元素的和为i，显然lookup[0]为真，其余位置初始化为假。最终我们要的是lookup[target]的真假，其中target为子数组和。那这样过程就很简单，只需要遍历每个元素，更新状态数组即可，状态更新的条件为：对于数i，如果i-n可以表示成原数组的若干元素的和，那么i也可以。下面是这个思路的代码：

def canPartition(self, nums: 'List[int]') -> 'bool':
    target, remainder = divmod(sum(nums), 2)
    if remainder is not 0:
        return False
    lookup = [True] + [False] * target
    for n in nums:
        for i in range(target, n - 1, -1):
            lookup[i] = lookup[i] or lookup[i - n]
    return lookup[target]

对于这个思路，还有一种巧妙的解法，那就是用bit maniputation的方法，在动态规划的题中也比较常见，即用比特位表示对应的数是否可以通过某些元素求和得到。举个例子，加入输入的数组nums是[2,3,4]，那么初始化status为1，遍历该数组，将status左移对应位得到的status为1的位置表示对应的数可以通过遍历过的元素求和得到。实现该方法的代码如下：

def canPartition2(self, nums: 'List[int]') -> 'bool':
    target, remainder = divmod(sum(nums), 2)
    if remainder is not 0:
        return False
    status = 1
    for n in nums:
        status = status << n | status
    return True if status >> target & 1 else False

partition to k equal subset sum

Given an array of integers nums and a positive integer k, find whether it’s possible to divide this array into k non-empty subsets whose sums are all equal.

Example 1:

Input: nums = [4, 3, 2, 3, 5, 2, 1], k = 4
Output: True
Explanation: It’s possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Note:

1 <= k <= len(nums) <= 16.
0 < nums[i] < 10000.

这道题是上面那道题的扩展，解法还是动态规划和回溯法，思想是一致的，不过要写出更一般化的代码。

首先是回溯法，我们可以从上道题的思路受到启发, 将groups的元素个数初始化为k个即可。对应的代码如下：


def canPartitionKSubsets(self, nums: 'List[int]', k: 'int') -> 'bool':
    # each subsets should sum to subsum
    subsum, remainder = divmod(sum(nums), k)
    if remainder or max(nums) > subsum:
        return False

    def search(groups):
        if not nums: return True
        v = nums.pop()
        for i, g in enumerate(groups):
            if g + v <= subsum:
                groups[i] += v
                if search(groups): return True
                groups[i] -= v
            if not g: break
        nums.append(v)
        return False 
    
    # remove the element in nums that equals to subset sum, just a trick to 
    # reduce time
    nums.sort()
    if nums[-1] > subsum: return False
    while nums and nums[-1] == subsum:
        k -= 1
        nums.pop()
    return search([0] * k)

值得注意的是，代码中有一句if not g: break，这句话的作用是用来减少循环次数的，具体如何减少，读者可以思考一下。

第二种方法就是动态规划了，鉴于这个思路leetcode上有清楚的解答，如果有看不懂的话可以留言，我再把中文解释加上。对应代码如下：


def canPartitionKSubsets2(self, nums: 'List[int]', k: 'int') -> 'bool':
    subsum, remainder = divmod(sum(nums), k)
    if remainder or max(nums) > subsum:
        return False
    mem = [None] * (1 << len(nums))
    mem[-1] = True
    def search(used, remain):
        if mem[used] is None:
            tar = (remain - 1) % subsum + 1
            mem[used] = any(search(used | (1 << i), remain - n) for i, n in
                    enumerate(nums) if (used >> i) & 1 == 0 and n <= tar)
        return mem[used]
    return search(0, subsum * k)

总结

这是两道经典的面试题，主要解法是用的比较多的动态规划和回溯法，而这两类解法都有比较固定的代码结构，只需要抓住关键条件对不同问题灵活运用即可。