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

Find the K-Sum of an Array

Number: 2462

Difficulty: Hard

Paid? No

Companies: Hubspot, Amazon


Problem Description

You are given an integer array nums and a positive integer k. A subsequence is an array derived from nums by deleting some or no elements while keeping the order of the remaining ones. The K-Sum of the array is defined as the kth largest subsequence sum (the empty subsequence sum is 0). Return the K-Sum of the array.


Key Insights

  • The largest subsequence sum is obtained by summing all positive numbers.
  • Negative numbers, when included in a subsequence, reduce the overall sum. We can view each negative number as a “penalty” with magnitude = |negative value|.
  • By computing posSum as the sum of all positive numbers, any subsequence sum can be written as: posSum minus the total penalty contributed by the negatives selected.
  • The problem reduces to generating the k smallest possible penalty sums from the list of absolute values of negative numbers.
  • A best-first search (using a min-heap) can be used to generate these penalty sums in increasing order without enumerating all (up to 2^(# of negatives)) possibilities.

Space and Time Complexity

Time Complexity: O(n log n + k log k)
 • O(n log n) to partition and sort the negative numbers’ absolute values.
 • O(k log k) in the worst-case for processing k elements from the heap.

Space Complexity: O(n + k)
 • O(n) for storing the penalties.
 • O(k) for the heap used in the best-first search.


Solution

  1. Calculate posSum as the sum of all positive numbers in nums.
  2. Collect the absolute values of all negative numbers (ignoring zeros since they do not change the sum) and sort them in ascending order.
  3. The goal is to determine penalty = (subset sum of these absolute values) such that the kth smallest penalty gives the kth largest subsequence sum (which equals posSum minus penalty).
  4. Use a min-heap (priority queue) to generate subset sums in increasing order. Initialize the heap with (currentPenalty = 0, currentIndex = 0).
  5. For each candidate (penalty, index) popped from the heap:
    • This candidate represents a penalty sum associated with a subsequence choice of negatives.
    • If index is less than the length of the penalty list, then generate two new candidates: • Include penaltyCandidates[index]: (penalty + penaltyCandidates[index], index + 1) • Exclude penaltyCandidates[index]: (penalty, index + 1)
    • Use a visited set to avoid processing the same state more than once.
  6. When you pop the kth candidate (starting count from 1) from the heap, the answer is posSum minus this penalty.
  7. Return the result.

Code Solutions

import heapq

def kSum(nums, k):
    # Calculate posSum (sum of all positive numbers)
    pos_sum = sum(x for x in nums if x > 0)
    
    # Get absolute values for negative numbers (ignore zeros)
    negatives = [abs(x) for x in nums if x < 0]
    negatives.sort()  # sorted in ascending order
    
    # Min-heap to generate increasing penalty values. Each element is (penalty, index)
    heap = [(0, 0)]
    visited = set((0, 0))  # to avoid duplicate states
    
    count = 0
    while heap:
        penalty, index = heapq.heappop(heap)
        count += 1
        # When we have popped k elements, we have the kth smallest penalty.
        if count == k:
            return pos_sum - penalty
        
        # If we can consider more negatives, generate next states.
        if index < len(negatives):
            # Option 1: Include the negative at current index (thus adding penalty)
            new_state = (penalty + negatives[index], index + 1)
            if new_state not in visited:
                visited.add(new_state)
                heapq.heappush(heap, new_state)
            # Option 2: Exclude the negative at current index
            new_state = (penalty, index + 1)
            if new_state not in visited:
                visited.add(new_state)
                heapq.heappush(heap, new_state)

# Example usage:
print(kSum([2, 4, -2], 5))  # Expected output: 2
← Back to All Questions