Append K Integers With Minimal Sum LeetCode solution in C++

1
class Solution {
2
 public:
3
  long long minimalKSum(vector<int> &nums, int k) {
4
    int n = nums.size();
5

6
    // Sort array so that we can get smaller elements first
7
    sort(begin(nums), end(nums));
8

9
    long long sum = 0;
10

11
    // Previous element that we encountered
12
    int prev = 0, curr = 0;
13

14
    for (int i = 0; i < n; i++) {
15
      curr = nums[i];
16

17
      // Get the difference b/w prev and current
18
      // So that we can append elements that are in between them
19
      long long diff = (curr - prev - 1);
20

21
      // If prev and current are same then just skip
22
      if (diff <= 0) {
23
        prev = curr;
24
        continue;
25
      }
26

27
      // If there are more available elements b/w prev and current
28
      // Then we just take closest k element from prev.
29
      // And leave remaining as it is
30
      if (diff > k) {
31
        diff = k;
32
        curr = prev + k + 1;
33
      }
34

35
      // Get the sum of all elements b/w prev and current
36
      // Since it is AP series, we can use direct formula
37
      sum += (diff * 1LL * (curr + prev) / 2);
38

39
      // Update previous to current
40
      prev = curr;
41

42
      // Update count of how many more element we need to append
43
      k -= diff;
44

45
      if (k == 0) break;
46
    }
47

48
    // Case : When we have reached the end of array
49
    // And still we have some more element left to append
50
    if (k) {
51
      sum += (k * 1LL * (2 * prev + k + 1) / 2);
52
    }
53

54
    return sum;
55
  }
56
};