Partition of a set into K subsets with equal sum

Posted 2020-10-16 apanda009

Given an integer array of N elements, the task is to divide this array into K non-empty subsets such that the sum of elements 

in every subset is same. All elements of this array should be part of exactly one partition.
Examples:

Input : arr = [2, 1, 4, 5, 6], K = 3
Output : Yes
we can divide above array into 3 parts with equal
sum as [[2, 4], [1, 5], [6]]

Input  : arr = [2, 1, 5, 5, 6], K = 3
Output : No
It is not possible to divide above array into 3
parts with equal sum

public class Test {
    boolean ans = false;
    public void sqrt(int[] nums, int k) {
        if (nums == null || nums.length < k || k < 0) {
            System.out.println(false);
            return;
        }
        if (k == 1) {
            return;
        }
        int sum = 0;
        for (int i : nums) {
            sum += i;
        }
        if (sum % k != 0) {
            return;
        }
        int ave = sum / k;
        int[] visited = new int[nums.length];

        Arrays.sort(nums);
        dfs(nums, k, visited, 0, 0, ave);
        System.out.println(ans);

    }
    public void dfs(int[] nums, int k, int[] visited,  int pos,int sum,int ave) {
        if (sum == ave) {
            if (pos == k - 1) {
                ans = true;
                return;
            }
            dfs(nums, k, visited, pos + 1, 0, ave);
        }
        if (sum > ave) {
            return;
        } else {
            for (int i = 0; i < nums.length; i++) {
                if (nums[i] + sum > ave || ans == true) {
                    return;
                }
                if (visited[i] != 1) {
                    sum += nums[i];
                    visited[i] = 1;
                    dfs(nums, k, visited, pos, sum, ave);

                    sum -= nums[i];
                    visited[i] = 0;
                }
            }
        }
    }