2386. Find the K-Sum of an Array

Problem Description

In this problem, you receive an array of integers called nums and a positive integer k. The goal is to figure out the kth largest sum of a subsequence within this array. A subsequence is a sequence that can be derived from the array by removing some or none of the items, while maintaining the order of the remaining items. Additionally, it's important to note that the empty subsequence (a sequence with no elements) has a sum of 0, and this should be considered when determining the K-Sum. The challenge then lies in efficiently finding the K-Sum because as the array grows with more elements, the number of subsequences that can be generated increases exponentially.

Intuition

To arrive at the solution, consider the nature of subsequences and the sums they generate. If the array contains positive numbers only, the largest subsequence sum is the sum of the entire array, and the smallest would be 0 (the empty subsequence). However, negative numbers complicate this because including a negative number in a subsequence can reduce the sum, so it may not always be beneficial to include all positive numbers in a subsequence.

The intuition behind the solution involves two key insights:

Max Sum with Positive Numbers: First, we calculate the maximum possible sum of the array as if all numbers were positive. This is because the highest K-Sum, specifically the 1st K-Sum, will not include any of the negative numbers (assuming we're only dealing with positive or non-positive numbers now).
A Priority Queue to Find the K-th Sum: After transforming all negative numbers into positive, we can sort the array. We then use a min-heap (priority queue) to systematically generate the sums of subsequences. By doing this, we ensure that we keep track of the smallest sums we can pop from the heap until we reach the kth smallest. We consider two operations with the heap: either we add the next number in the sorted array to the current sum (expanding the subsequence by one), or we replace the last added number with the next number (changing the last element of the subsequence), but only if it's not the first element in the subsequence (to avoid duplicates).

These steps allow us to cleverly sidestep the brute-force approach of generating every possible subsequence and its sum, which would be impractical for larger arrays.

Learn more about Sorting and Heap (Priority Queue) patterns.

Solution Approach

The solution approach utilizes several algorithms and data structures to efficiently find the K-Sum of the array.

Dealing with Max Sum: We initialize a variable mx with 0 to store the maximum sum. We iterate over nums and check each value. If the value is positive, we add it to mx. If it is negative (or non-positive, to be more general), it means including this value in our subsequence would decrease the sum. So, we take the absolute value of that number and then include it in the array nums, effectively transforming the array to contain only non-negative numbers. Later, we sort this transformed array, which is crucial for the priority queue operations.
Priority Queue (Min-Heap): A priority queue is implemented using a min-heap to keep track of the smallest sums that could lead to the K-Sum when combined with the subsequent elements in the array. We use this to efficiently get the kth smallest sum. The heap is initialized with a tuple containing 0 (the sum of the empty subsequence) and 0 (the index indicating we haven't started to use any number from nums).
Heap Operation to Find K-th Largest Sum: For each k-1 iterations (since we're trying to find the kth largest, we do one less), we pop the smallest sum from the heap. For each popped sum, which comes as a tuple (sum, index), we check if index is less than the length of nums. If it is, we push a new tuple onto the heap: the current sum s plus the number at the current index nums[i], and increment the index by 1. This operation represents expanding the subsequence to include the nums[i].
Avoiding Duplication: If the index i is not zero, we push yet another tuple onto the heap which accounts for swapping the last number of the subsequence with the current number at index i of nums. The calculation is s + nums[i] - nums[i-1], and the index is incremented by 1. This operation is crucial because it effectively generates sums that include the current number, but not the previous, preventing overcounting and ensuring that all subsequences considered are unique.
Retrieve the K-th Largest Sum: After performing the for-loop of k-1 iterations, the heap's smallest sum is the kth largest subsequence sum because all smaller sums have been popped off the heap already.
Calculating the Final Answer: Since mx represents the largest sum we could possibly make without negative numbers, and h[0][0] contains the smallest sum at the kth position, subtracting h[0][0] from mx gives us the actual value of the kth largest subsequence sum, even when considering the original array with its negative numbers.

In summary, the implementation cleverly leverages a heap to simulate the process of generating subsequence sums and smartly navigates the subsequence space to find the K-Sum efficiently without generating all possible subsequences.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's use a small example to illustrate the solution approach.

Consider the array nums = [3, -1, 2] and k = 2. The task is to find the 2nd largest sum of a subsequence within this array.

Dealing with Max Sum: First, we initialize mx with 0. After iterating over nums, we will transform it into [3, 1, 2] as we convert the negative number to a positive by taking the absolute value. The mx will now be the sum of this array, which is 6.
Priority Queue (Min-Heap): Initialize a min-heap that starts with the tuple (0, 0), indicating the sum of the empty subsequence and the starting index.
Heap Operation to Find 2nd Largest Sum: Now, we will perform k - 1 iterations over the heap to find the 2nd largest sum:
- Pop the top of the heap which is (0, 0) and since index is less than the length of nums, push (0 + nums[0], 0 + 1) which is (3, 1) and (0 + nums[1], 1 + 1) which is (1, 2) into the heap.
- Now the heap contains (1, 2) and (3, 1). We pop the smallest sum again, which is (1, 2), and since index is less than the length of nums, we push (1 + nums[2], 2 + 1) which is (3, 3) into the heap.
Avoiding Duplication: There isn't any duplication avoidance necessary in this example since there are no numbers in nums to follow nums[2].
Retrieve the 2nd Largest Sum: After k - 1 iterations, our heap contains (3, 1) and (3, 3). The smallest sum at the top of the heap is (3, 1), which represents the 2nd largest subsequence sum.
Calculating the Final Answer: mx is 6, and the kth smallest sum from our heap is 3. So, the 2nd largest subsequence sum is mx - heap's top sum, which equals 6 - 3 = 3.

The 2nd largest sum of a subsequence in the array [3, -1, 2] is 3. This sum can be obtained from the subsequence [3] or [2, -1, 3] in the original array before we transformed the negative numbers into positive.

Solution Implementation

1from heapq import heappush, heappop
2
3class Solution:
4    def kthLargestSum(self, nums: List[int], k: int) -> int:
5        # Calculate the sum of positive elements, and make all elements positive
6        max_sum = 0
7        for index, value in enumerate(nums):
8            if value > 0:
9                max_sum += value
10            else:
11                nums[index] = -value
12        nums.sort()  # Sort the array
13      
14        # Create a min-heap to store the sums and their corresponding indexes
15        min_heap = [(0, 0)]
16      
17        # Pop and push the next two sums to the heap (k - 1) times
18        for _ in range(k - 1):
19            current_sum, index = heappop(min_heap)
20            if index < len(nums):
21                # Push the next sum which includes the nums[index]
22                heappush(min_heap, (current_sum + nums[index], index + 1))
23              
24                # Avoid duplicating the first element of the sorted nums array
25                if index:
26                    # Calculate the sum by swapping out the previous element (thus maintaining the count of elements in current sum)
27                    heappush(min_heap, (current_sum + nums[index] - nums[index - 1], index + 1))
28      
29        # The k-th largest sum is the max_sum minus the smallest sum in the heap
30        return max_sum - min_heap[0][0]
31
32# Note: The code assumes the existence of `List` imported from `typing` module.
33# If not present, add: from typing import List
34

1import java.util.PriorityQueue;
2import java.util.Comparator;
3import java.util.Arrays;
4import javafx.util.Pair; // Make sure to import the correct package for the Pair class based on the development environment
5
6class Solution {
7    public long kSum(int[] nums, int k) {
8        // Initialize variable to store the sum of positive numbers.
9        long maxSumPositives = 0;
10        int length = nums.length;
11      
12        // Convert negative numbers to positive and sum all positive numbers.
13        for (int i = 0; i < length; ++i) {
14            if (nums[i] > 0) {
15                maxSumPositives += nums[i];
16            } else {
17                nums[i] = -nums[i];
18            }
19        }
20      
21        // Sort the modified array in non-decreasing order.
22        Arrays.sort(nums);
23      
24        // Use a priority queue to keep track of the sum-pairs efficiently.
25        PriorityQueue<Pair<Long, Integer>> minHeap = new PriorityQueue<>(Comparator.comparing(Pair::getKey));
26        minHeap.offer(new Pair<>(0L, 0)); // Offer initial pair of (sum=0, index=0).
27      
28        // Loop until we find the kth smallest sum.
29        while (--k > 0) {
30            Pair<Long, Integer> currentPair = minHeap.poll(); // Poll the smallest sum pair.
31            long currentSum = currentPair.getKey();
32            int currentIndex = currentPair.getValue();
33          
34            // If there is a next index, offer new pairs into the priority queue.
35            if (currentIndex < length) {
36                minHeap.offer(new Pair<>(currentSum + nums[currentIndex], currentIndex + 1));
37                // Avoid duplicate sums by checking if currentIndex is greater than 0.
38                if (currentIndex > 0) {
39                    minHeap.offer(new Pair<>(currentSum + nums[currentIndex] - nums[currentIndex - 1], currentIndex + 1));
40                }
41            }
42        }
43      
44        // Return the maximum sum of positives minus the k-th smallest sum (the top element in the queue).
45        return maxSumPositives - minHeap.peek().getKey();
46    }
47}
48

1#include <vector>
2#include <queue>
3#include <algorithm>
4using namespace std;
5
6// Define pair type with long long and int for use in priority queue
7using PairLLInt = pair<long long, int>;
8
9class Solution {
10public:
11    long long kSum(vector<int>& nums, int k) {
12        long long maxSumNonNegatives = 0; // To store the sum of non-negative numbers
13        int size = nums.size();
14        // Convert negative numbers to positive and calculate maxSumNonNegatives
15        for (int i = 0; i < size; ++i) {
16            if (nums[i] > 0) {
17                maxSumNonNegatives += nums[i];
18            } else {
19                nums[i] = -nums[i];
20            }
21        }
22        // Sort nums to utilize the non-decreasing property in generating sums
23        sort(nums.begin(), nums.end());
24      
25        // Min-heap to store the smallest sums and their indices
26        priority_queue<PairLLInt, vector<PairLLInt>, greater<PairLLInt>> minHeap;
27        minHeap.push({0, 0}); // Initial pair
28
29        // Generate the first k-1 sums
30        while (--k) {
31            auto [sum, idx] = minHeap.top();
32            minHeap.pop();
33
34            if (idx < size) {
35                minHeap.push({sum + nums[idx], idx + 1});
36                // Avoid duplicates by checking if the current sum results from
37                // a combination of previous numbers
38                if (idx > 0) {
39                    minHeap.push({sum + nums[idx] - nums[idx - 1], idx + 1});
40                }
41            }
42        }
43
44        // k-th sum is the smallest sum which is the difference
45        // between maxSumNonNegatives and the top element of min-heap
46        return maxSumNonNegatives - minHeap.top().first;
47    }
48};
49

1function kSum(nums: number[], k: number): number {
2    let maxSumNonNegatives: number = 0; // Store the sum of non-negative numbers
3    let size: number = nums.length;
4
5    // Convert negative numbers to positive and calculate maxSumNonNegatives
6    for (let i = 0; i < size; ++i) {
7        if (nums[i] > 0) {
8            maxSumNonNegatives += nums[i];
9        } else {
10            nums[i] = -nums[i];
11        }
12    }
13
14    // Sort nums to utilize its non-decreasing property in generating sums
15    nums.sort((a, b) => a - b);
16
17    // Define the priority queue to store the smallest sums and their indices
18    let minHeap: Array<[number, number]> = [];
19
20    // Function to add elements to the min-heap
21    const addToMinHeap = (sum: number, index: number): void => {
22        minHeap.push([sum, index]);
23        // Percolate up to maintain heap invariant
24        let current = minHeap.length - 1;
25        while (current > 0) {
26            let parent = Math.floor((current - 1) / 2);
27            if (minHeap[current][0] < minHeap[parent][0]) {
28                [minHeap[current], minHeap[parent]] = [minHeap[parent], minHeap[current]];
29                current = parent;
30            } else {
31                break;
32            }
33        }
34    };
35
36    // Function to extract the smallest element from the min-heap
37    const extractFromMinHeap = (): [number, number] => {
38        const smallest = minHeap[0];
39        const lastElement = minHeap.pop()!;
40        if (minHeap.length > 0) {
41            minHeap[0] = lastElement;
42            // Percolate down to maintain heap invariant
43            let current = 0;
44            while (true) {
45                let leftChild = current * 2 + 1;
46                let rightChild = current * 2 + 2;
47                let smallestChild = current;
48              
49                if (leftChild < minHeap.length && minHeap[leftChild][0] < minHeap[smallestChild][0]) {
50                    smallestChild = leftChild;
51                }
52                if (rightChild < minHeap.length && minHeap[rightChild][0] < minHeap[smallestChild][0]) {
53                    smallestChild = rightChild;
54                }
55                if (smallestChild !== current) {
56                    [minHeap[current], minHeap[smallestChild]] = [minHeap[smallestChild], minHeap[current]];
57                    current = smallestChild;
58                } else {
59                    break;
60                }
61            }
62        }
63        return smallest;
64    };
65
66    addToMinHeap(0, 0); // Initial element
67
68    // Generate the first k-1 sums
69    while (--k) {
70        const [sum, idx] = extractFromMinHeap();
71
72        if (idx < size) {
73            addToMinHeap(sum + nums[idx], idx + 1);
74            // Avoid duplicates by checking if the current sum is from
75            // a combination of previous numbers
76            if (idx > 0) {
77                addToMinHeap(sum + nums[idx] - nums[idx - 1], idx + 1);
78            }
79        }
80    }
81
82    // k-th sum is the smallest sum which is the difference
83    // between maxSumNonNegatives and the top element of min-heap
84    const [smallestSum, ] = extractFromMinHeap();
85    return maxSumNonNegatives - smallestSum;
86}
87
88// Usage example:
89// const nums: number[] = [1, 2, 3];
90// const k: number = 2;
91// console.log(kSum(nums, k)); // Call the kSum function with the example parameters
92

Time and Space Complexity

Time Complexity

The time complexity of the kSum function involves several components:

Initialization and Sorting: Initializing mx and converting negative numbers in the nums list to positive ones takes O(n) time where n is the length of nums. Sorting the modified nums list takes O(n log n) time.
Heap Operations: The function performs a series of heap operations, namely heappop and heappush, within a loop that runs k - 1 times. Each heap operation can be done in O(log k) time. However, in the worst-case scenario, there can be up to 2*(k-1) elements in the heap, as for each pop operation, potentially two elements are pushed back. This leads to a time complexity of O((k - 1) * log(k)) for all heap operations.

Combining these operations, the overall time complexity is determined by O(n log n + (k - 1) * log(k)). Since the heap operations depend on k, the larger of n log n and (k - 1) * log(k) will dominate the time complexity.

Thus, the total time complexity is O(max(n log n, (k - 1) * log(k))).

Space Complexity

The space complexity of the kSum function consists of:

Heap Space: The heap can contain up to 2*(k-1) elements in the worst-case scenario, which contributes O(k) space complexity.
Other Variables: The space used by variables mx, s, and i, is O(1).

Therefore, the total space complexity of the function is O(k) since this is the term that grows with the input size and dominates the space complexity.

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

What's the output of running the following function using input [30, 20, 10, 100, 33, 12]?

1def fun(arr: List[int]) -> List[int]:
2    import heapq
3    heapq.heapify(arr)
4    res = []
5    for i in range(3):
6        res.append(heapq.heappop(arr))
7    return res
8

1public static int[] fun(int[] arr) {
2    int[] res = new int[3];
3    PriorityQueue<Integer> heap = new PriorityQueue<>();
4    for (int i = 0; i < arr.length; i++) {
5        heap.add(arr[i]);
6    }
7    for (int i = 0; i < 3; i++) {
8        res[i] = heap.poll();
9    }
10    return res;
11}
12

1class HeapItem {
2    constructor(item, priority = item) {
3        this.item = item;
4        this.priority = priority;
5    }
6}
7
8class MinHeap {
9    constructor() {
10        this.heap = [];
11    }
12
13    push(node) {
14        // insert the new node at the end of the heap array
15        this.heap.push(node);
16        // find the correct position for the new node
17        this.bubble_up();
18    }
19
20    bubble_up() {
21        let index = this.heap.length - 1;
22
23        while (index > 0) {
24            const element = this.heap[index];
25            const parentIndex = Math.floor((index - 1) / 2);
26            const parent = this.heap[parentIndex];
27
28            if (parent.priority <= element.priority) break;
29            // if the parent is bigger than the child then swap the parent and child
30            this.heap[index] = parent;
31            this.heap[parentIndex] = element;
32            index = parentIndex;
33        }
34    }
35
36    pop() {
37        const min = this.heap[0];
38        this.heap[0] = this.heap[this.size() - 1];
39        this.heap.pop();
40        this.bubble_down();
41        return min;
42    }
43
44    bubble_down() {
45        let index = 0;
46        let min = index;
47        const n = this.heap.length;
48
49        while (index < n) {
50            const left = 2 * index + 1;
51            const right = left + 1;
52
53            if (left < n && this.heap[left].priority < this.heap[min].priority) {
54                min = left;
55            }
56            if (right < n && this.heap[right].priority < this.heap[min].priority) {
57                min = right;
58            }
59            if (min === index) break;
60            [this.heap[min], this.heap[index]] = [this.heap[index], this.heap[min]];
61            index = min;
62        }
63    }
64
65    peek() {
66        return this.heap[0];
67    }
68
69    size() {
70        return this.heap.length;
71    }
72}
73
74function fun(arr) {
75    const heap = new MinHeap();
76    for (const x of arr) {
77        heap.push(new HeapItem(x));
78    }
79    const res = [];
80    for (let i = 0; i < 3; i++) {
81        res.push(heap.pop().item);
82    }
83    return res;
84}
85

[30, 20, 10]

[100, 33, 30]

[10, 12, 20]

I don’t know