LeetCode算法问题12 —— Partition to K Equal Sum Subsets

看一下问题描述

问题要求判断一个给定的数组能否划分为k个代数和相同的子数组。

首先必须获取给定数组的和

 int sum = 0;
 for(int num:nums)
     sum += num;

问题就变成了将每个元素划分到哪个子数组的问题了，我的想法是先让第一个子数组从起始位置开始挑选数字看能否组成一个和为sum/k的子数组。
那么这样子就要创建k个子数组了吗？答案是否定的，我创建了一个数组visited标识给定数组的每个位i，如果它已被划分进入一个子数组，则将其visited[i] = 1，否则置为0。同时封装一个递归函数
bool canPartition(vector& nums, vector& visited, int start_index, int number, int cur_sum, int cur_num, int target)
它用于找从nums的第start_index位开始为第k-number个数组找和为target的数组，其中当前数组已有cur_num个数且和为cur_sum。
对于这一个递归函数，当我们找到了这个子数组的所有数后，执行如下代码

if(cur_sum == target && cur_num >0 )return canPartition(nums, visited, 0, number-1, 0, 0, target);

意为开始查找下一个子数组的元素，即第k-number+1个数组
如果还没找满，则start_index向后遍历，对于还未放入数组的数执行

if(!visited[i]){
                visited[i] = 1;
                if(canPartition(nums, visited, i+1, number, cur_sum + nums[i], cur_num++, target))return true;
                visited[i] = 0;
            }

意为如果将这个数放入数组，能否找到一个和为target的数组。
而整个大结构是这样的

for(int i = start_index; i<nums.size(); i++){
            if(!visited[i]){
                visited[i] = 1;
                if(canPartition(nums, visited, i+1, number, cur_sum + nums[i], cur_num++, target))
                    return true;
                visited[i] = 0;
            }
        }

这里有回溯的思想。

由此即可解决该问题。

完整代码如下

class Solution {
public:
    bool canPartitionKSubsets(vector<int>& nums, int k) {
        int sum = 0;
        for(int num:nums)
            sum += num;
        if(k <= 0 || sum % k != 0)
            return false;
        vector<int> visited(nums.size(), 0);
        return canPartition(nums, visited, 0, k, 0, 0, sum/k);
    }

    bool canPartition(vector<int>& nums, vector<int>& visited, int start_index, int number, int cur_sum, int cur_num, int target){
        if(number == 1)
            return true;
        if(cur_sum == target && cur_num >0 )return canPartition(nums, visited, 0, number-1, 0, 0, target);
        for(int i = start_index; i<nums.size(); i++){
            if(!visited[i]){
                visited[i] = 1;
                if(canPartition(nums, visited, i+1, number, cur_sum + nums[i], cur_num++, target))return true;
                visited[i] = 0;
            }
        }
        return false;
    }
};