We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Number of Great Partitions

Number: 2601

Difficulty: Hard

Paid? No

Companies: Sprinklr

Problem Description

Given an array of positive integers and an integer k, partition the array into two ordered groups (each element must go to exactly one group) so that the sum of elements in each group is at least k. Two partitions are distinct if there exists at least one index for which the element is assigned to a different group. Return the number of distinct “great” partitions modulo 10^9 + 7.

Key Insights

  • There are 2^n total ways to assign each element to one of the two groups (order matters).
  • A partition is "great" if the sum of elements in both groups is at least k; equivalently, if one group has sum S, then S must be between k and (totalSum - k) inclusive.
  • It is easier to count the number of partitions that fail the condition (i.e. where one group’s sum is less than k) and subtract this from the total.
  • Use dynamic programming over subset sums but note that we only need to count sums less than k.
  • Beware of double counting: partitions that are invalid for both groups (which can occur only when totalSum < 2k) need to be added back using inclusion–exclusion.

Space and Time Complexity

Time Complexity: O(n * k)
Space Complexity: O(k)

Solution

We first note that every partition corresponds to choosing a subset of elements for the first group (and the rest go to the second group). The “great” partition condition becomes: let S1 be the sum of the first group; then S1 >= k and S2 = totalSum - S1 >= k. This implies S1 must be in the range [k, totalSum - k].

We can compute the number of ways to assign a subset of nums such that its sum is less than k using dynamic programming. Here, dp[s] represents the number of ways to achieve a sum s (for group one) using some subset of elements, with the constraint that we only care about s < k. For each number in the array, we update dp by:

  • Keeping the number out of group one.
  • Adding the number to group one if doing so does not reach k (i.e. if the new sum is < k).

Let f = sum{dp[s]} for 0 <= s < k. This f represents the number of assignments where group one has sum less than k. By symmetry, the number of assignments where group two’s sum is less than k is also f. However, assignments where both groups have sum < k (which can happen only if totalSum < 2k) are counted twice. Denote the number of ways that both groups fail as g. They are counted by considering only those assignments where S1 is in the range [totalSum - k + 1, k - 1].

Finally, the answer is computed as:
Answer = 2^n - (f + f - g), taken modulo 10^9 + 7.

This technique avoids iterating over all possible partitions by leveraging subset sum DP up to k, which is sufficient since k <= 1000.

Code Solutions

MOD = 10**9 + 7

def numberOfGreatPartitions(nums, k):
    n = len(nums)
    totalSum = sum(nums)
    # Compute dp[s] = ways to get sum s for group1 (only s < k considered)
    dp = [0] * k
    dp[0] = 1
    for num in nums:
        # For each num, update dp in reverse order so that you do not count the same number twice.
        new_dp = dp[:]  # copying current state (not choosing num -> group1)
        for s in range(k):
            if dp[s] and s + num < k:
                new_dp[s + num] = (new_dp[s + num] + dp[s]) % MOD
        dp = new_dp
    
    # Count assignments where group1 sum < k
    invalid_group1 = sum(dp) % MOD

    # Count partitions invalid for both groups.
    both_invalid = 0
    # Only possible if totalSum < 2*k because otherwise at least one group must have sum >= k.
    if totalSum < 2 * k:
        # For any valid assignment with group1 sum = s, group2 sum = totalSum - s < k if and only if s > totalSum - k.
        low_bound = max(totalSum - k + 1, 0)
        for s in range(low_bound, k):
            both_invalid = (both_invalid + dp[s]) % MOD

    total_assignments = pow(2, n, MOD)
    # Remove assignments where one of the groups is invalid (and add back those subtracted twice).
    result = (total_assignments - (invalid_group1 + invalid_group1 - both_invalid)) % MOD
    return result

# Example usage:
print(numberOfGreatPartitions([1,2,3,4], 4))  # Expected output: 6
print(numberOfGreatPartitions([3,3,3], 4))    # Expected output: 0
print(numberOfGreatPartitions([6,6], 2))      # Expected output: 2
← Back to All Questions