Author: Nishant Gupta
Tester: Priyanshu Srivastava
Editorialist: Nishant Gupta
Meet in middle
Given N numbers and three integers K, A, B you have to tell the number of subsets of size at most K and whose Niceness value is in the range [A, B]. Niceness value of a number is product of its prime divisors and Niceness value of a subset is sum of niceness values of its elements.
Niceness value of each number is calculated and stored in another array. Split the niceness array in two halves, and generate their subsets independently. Store the sum of subsets in a 2-D vector like sum of subset of size x will be stored in vec[x]. Sort the subsets of second half of the array with respect to the sum of the subset for each size. For every subset of size (say s) ans sum (say sum) in first half, do a binary search for A-sum (lower_bound) and B-sum (upper bound) for every size, [0…MIN(N/2, K-s)]. Add the difference of upper bound and lower bound for each subset of first half.
Niceness value can be calculated by iterating over each element in the array and then with O(sqrt(a[i])) complexity all the prime factors can be found and hence their product. Store this in array niceness. Now question reduces to finding subsets of niceness with size at most K and sum in the range [A, B].
Now there are at most 32 numbers and if one has to calculate number of subsets whose sum is in the range [A, B], it will reach to approximately 232 * 32 iterations in worst case which is not feasible. So one has to split the niceness array into two halves firsthalf, 0…mid and secondhalf, mid+1…N and generate their subsets independently. This would be taking at most 216 * 16 iterations which is feasible. Store the sum of a subset in 2-D vector, vec[i] will contain sum of all subsets of size i. Sort the subsets with respect to their sum for each size i in range 0…N/2.
Now for every subset of firsthalf of size = i and sum = SUM, we have to count the subsets of secondhalf with size 0…MIN(N/2, K-i] such that maximum size of subset remains K. To count the subsets such that sum is in range [A, B], take lower bound for A-SUM and upper bound for B-SUM . Take the difference of these two values and this will give you count of subsets of secondhalf who will sum up with subset of firsthalf with sum = SUM and will give total sum in the range [A, B]. This is done for all subsets of firsthalf and total subsets thus calculated gives the answer.