Skip to content

2386. Find the K-Sum of an Array 👍

  • Time: $O(\max(\texttt{sort}(\texttt{nums}), k\log k))$
  • Space: $O(n + k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution {
 public:
  long long kSum(vector<int>& nums, int k) {
    const long long maxSum = getMaxSum(nums);
    const vector<int> absNums = getAbsNums(nums);
    long long ans = maxSum;
    // (the next maximum sum, the next index i)
    using P = pair<long long, int>;
    priority_queue<P> maxHeap;
    maxHeap.emplace(maxSum - absNums[0], 0);

    for (int j = 0; j < k - 1; ++j) {
      const auto [nextMaxSum, i] = maxHeap.top();
      maxHeap.pop();
      ans = nextMaxSum;
      if (i + 1 < absNums.size()) {
        maxHeap.emplace(nextMaxSum - absNums[i + 1], i + 1);
        maxHeap.emplace(nextMaxSum - absNums[i + 1] + absNums[i], i + 1);
      }
    }

    return ans;
  }

 private:
  long long getMaxSum(const vector<int>& nums) {
    long long maxSum = 0;
    for (const int num : nums)
      if (num > 0)
        maxSum += num;
    return maxSum;
  }

  vector<int> getAbsNums(const vector<int>& nums) {
    vector<int> absNums;
    for (const int num : nums)
      absNums.push_back(abs(num));
    ranges::sort(absNums);
    return absNums;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Solution {
  public long kSum(int[] nums, int k) {
    final long maxSum = getMaxSum(nums);
    final int[] absNums = getAbsNums(nums);
    long ans = maxSum;
    // (the next maximum sum, the next index i)
    Queue<Pair<Long, Integer>> maxHeap =
        new PriorityQueue<>((a, b) -> Long.compare(b.getKey(), a.getKey()));
    maxHeap.offer(new Pair<>(maxSum - absNums[0], 0));

    for (int j = 0; j < k - 1; ++j) {
      Pair<Long, Integer> pair = maxHeap.poll();
      final long nextMaxSum = pair.getKey();
      final int i = pair.getValue();
      ans = nextMaxSum;
      if (i + 1 < absNums.length) {
        maxHeap.offer(new Pair<>(nextMaxSum - absNums[i + 1], i + 1));
        maxHeap.offer(new Pair<>(nextMaxSum - absNums[i + 1] + absNums[i], i + 1));
      }
    }

    return ans;
  }

  private long getMaxSum(int[] nums) {
    long maxSum = 0;
    for (final int num : nums)
      if (num > 0)
        maxSum += num;
    return maxSum;
  }

  private int[] getAbsNums(int[] nums) {
    for (int i = 0; i < nums.length; ++i)
      nums[i] = Math.abs(nums[i]);
    Arrays.sort(nums);
    return nums;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution:
  def kSum(self, nums: List[int], k: int) -> int:
    maxSum = sum(num for num in nums if num > 0)
    absNums = sorted(abs(num) for num in nums)
    # (the next maximum sum, the next index i)
    maxHeap = [(-(maxSum - absNums[0]), 0)]
    nextMaxSum = maxSum

    for _ in range(k - 1):
      nextMaxSum, i = heapq.heappop(maxHeap)
      nextMaxSum *= -1
      if i + 1 < len(absNums):
        heapq.heappush(maxHeap, (-(nextMaxSum - absNums[i + 1]), i + 1))
        heapq.heappush(
            maxHeap, (-(nextMaxSum - absNums[i + 1] + absNums[i]), i + 1))

    return nextMaxSum