Leetcode 698 partition to K equal sum subsets (DFS pruning)

Given an integer array nums and an integer k, return true if it is 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 is possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.

Example 2:

Input: nums = [1,2,3,4], k = 3
Output: false

Constraints:

• 1 <= k <= nums.length <= 16
• 1 <= nums[i] <= 104
• The frequency of each element is in the range [1, 4].

Topic link ：https://leetcode.com/problems/partition-to-k-equal-sum-subsets/

The main idea of the topic ：n A number asks if it can be divided into k And equal portions

Topic analysis ： Classic search plus pruning problem , Two situations that definitely don't work ：1. Sum is not k Multiple ,2. The maximum value is larger than the mean value . Before searching, you need to sort the array from large to small , Because numbers with large values are less flexible , It is faster to encounter a situation without a solution than a small number , Because the subject restriction clearly indicates that the frequency of each number is 1 To 4, When it comes to i When making a decision with a number , If i-1 No choice , And nums[i] and nums[i - 1] equal , The first i A position must not be chosen

4ms, Time beats 75.5%

class Solution {

    public boolean dfs(int pos, int curSum, int each, int curSz, int k, int[] nums, boolean[] vis) {
        if (curSz == k) {
            return true;
        }
        for (int i = pos; i < nums.length; i++) {
            if (i > 0 && !vis[i - 1] && nums[i] == nums[i - 1]) {
                continue;
            }
            if (!vis[i] && curSum + nums[i] <= each) {
                vis[i] = true;
                if (curSum + nums[i] == each) {
                    if (dfs(0, 0, each, curSz + 1, k, nums, vis)) {
                        return true;
                    }
                } else {
                    if (dfs(i + 1, curSum + nums[i], each, curSz, k, nums, vis)) {
                        return true;
                    }
                }
                vis[i] = false;
            }
        }
        return false;
    }
    
    public boolean canPartitionKSubsets(int[] nums, int k) {
        Arrays.sort(nums);
        int sum = 0, n = nums.length;
        for (int i = 0; i < n / 2; i++) {
            int tmp = nums[i];
            nums[i] = nums[n - i - 1];
            nums[n - i - 1] = tmp;
        }
        for (int num : nums) {
            sum += num;
        }
        if (sum % k != 0) {
            return false;
        }
        int each = sum / k;
        if (nums[0] > each) {
            return false;
        }
        boolean[] vis = new boolean[n];
        return dfs(0, 0, each, 0, k, nums, vis);
    }
}