PROBLEM LINK:
Author: Jay Pandya
Tester 1: Minako Kojima
Tester 2: Shiplu Hawlader
Editorialist: Pawel Kacprzak
DIFFICULTY:
EASY
PREREQUISITES:
DP, bits
PROBLEM:
Given a set S of N integers the task is decide if it is possible to divide them into K non-empty subsets such that the sum of elements in every of the K subsets is equal.
QUICK EXPLANATION:
We use dynamic programming to solve it. If the solution exists, each subsets has to have a sum of its elements equal to the sum of all elements in S divided by K. Let X be that value. Let dp[k][bitmask] = 1 if and only if it is possible to divide a subset A of S denoted by the bitmask into k subsets each with sum X or to divide A into k - 1 subsets of sum X and one subset of sum < X. For a single dp[k][bitmask] entry, we will iterate over all elements of S which are not in A and try to extend current solution. The answer is “yes” if and only if dp[K][bitmask denoting the whole set S] = 1, and because we try to extend dp states only for k < K, we are sure than dp[K][bitmask] denotes if it is possible to divide a subset denoted by bitmask into K subsets with equal sums (see explanation below).
EXPLANATION:
Let SUM be the sum of all elements in S. If SUM % K != 0, then the answer is “no” because there is no way to divide elements equally.
Let X = SUM / K.
Let dp[k][bitmask] = 1 if and only if it is possible to divide a subset A of S denoted by the bitmask into k subsets each with sum X or to divide A into k - 1 subsets with sum X and one subset with sum < X.
Initially all entries in dp are set to 0.
At the beginning, we set dp[0][0] = 1 because it is possible to create 0 subsets with equal sums using no elements.
In the main loop, we iterate over all values k = 0, 1, …, K - 1 denoting the number of subsets for which we try to extend our solution and over all subsets of S denoted by a bitmask. Let A be a subset denoted by a bitmask. For a given dp[k][bitmask] we compute how much we have to add to the k+1th subset in order to get k+1 subsets, each with a sum X. We do that by subtracting k * X from the sum of elements in A. Then we iterate over all elements of S which are not in A and we try to extend our solution (see the code below):
Let a be an array denoting the set S. I omit the case when SUM % K != 0.
for k = 0 to K - 1:
for bitmask = 0 to 2^n - 1:
if not dp[k][bitmask]:
continue
sum = 0
for i = 0 to n - 1:
if (bitmask & (1LL << i)):
sum += a[i]
sum -= k * x
for i = 0 to n - 1:
if (bitmask & (1LL << i)):
continue //there is nothing to extend
new_mask = bitmask | (1LL << i) //bitmask denoting a set with i-th element added
if sum + a[i] == x: //we can fill the k+1 th subset with elements of sum X using a set of elements denoted by new_mask
dp[k + 1][new_mask] = 1
else if sum + a[i] < x: //we can fill k subsets with elements of sum X and one subsets with sum < X using a set of elements denoted by new_mask
dp[k][new_mask] = 1
if dp[K][2^n - 1] == 1:
print "yes"
else:
print "no"
Time Complexity:
The time complexity per one testcase is O(K * 2^N * N)
SOLUTIONS:
Author’s solution can be found here.
Tester’s solution can be found here.
Editorialist’s solution can be found here.
RELATED PROBLEMS:
To be uploaded soon.
