This problem looks busy at first, but the accepted solution is built around one steady invariant. For 658. Find K Closest Elements, the solution in this repository is mainly a two pointers solution.
Guide
What?
We want to turn the problem statement into a smaller set of decisions the computer can repeat safely. Instead of trying to be clever immediately, read the code as a sequence of questions:
- What state are we keeping?
- How do we move from one state to the next?
- When do we know the answer is already determined?
For this file, the main tools are: two pointers, sliding window, greedy.
Guide
When?
This is the kind of solution you want when the problem has structure hiding inside a messy-looking input. The accepted code reduces that pressure by storing exactly the information that remains useful later.
The important function names to track are findClosestElements.
Guide
Why?
The win comes from making each line carry responsibility: store the useful state, discard the rest, keep moving.
- Sorting is used to make local choices comparable, so the later scan does not have to rediscover order.
- The final return is not magic; it is the invariant after the loops or recursion have finished doing their accounting.
Guide
How?
Walk through the solution in this order:
- Start from the smallest reliable state.
- Expand one legal move at a time.
- Cache, count, or merge information as soon as it becomes settled.
- Let the final stored value answer the original question.
The most important competitive-programming habit here is to trust the invariant. Once the invariant is right, the loops become much less scary.
Guide
Complexity
- Time: O(n) to O(n log n), depending on the dominant loop or data structure operation
- Space: O(n) in the usual case for auxiliary containers or recursion
Guide
C++ Solution
Your submission
The accepted solution
solution.cpp
01//Runtime: 72 ms, faster than 99.56% of C++ online submissions for Find K Closest Elements.
02//Memory Usage: 31 MB, less than 16.67% of C++ online submissions for Find K Closest Elements.
03class Solution {
04public:
05 vector<int> findClosestElements(vector<int>& arr, int k, int x) {
06 auto it = lower_bound(arr.begin(), arr.end(), x);
07
08 if(it == arr.begin()){
09 return vector<int>(arr.begin(), arr.begin()+k);
10 }else if(it == arr.end()){
11 return vector<int>(arr.end() - k, arr.end());
12 }
13
14 int n = arr.size();
15 int idx = it - arr.begin();
16 int left, right;
17 left = idx - k/2;
18 right = left + k;
19
20 // cout << left << ", " << right << " -> ";
21
22 if(left < 0){
23 left = 0;
24 right = left +k;
25 }
26
27 if(right > n){
28 right = n;
29 left = right - k;
30 }
31
32 // cout << left << ", " << right << " -> ";
33
34 /*
35 since arr[idx] is just a surrogate of x,
36 we need to examine the diff of elements with "x" again
37 */
38
39 /*
40 while the diff of the element to the leftmost element "<=" that of the rightmost element,
41 move the subarray left
42
43 here the "<=" means we prefer smaller element
44 */
45 while(left > 0 && x - arr[left-1] <= arr[right-1] - x){
46 right--;
47 left--;
48 }
49
50 /*
51 while the diff of rightmost element "<" that of the element to the leftmost element,
52 move the subarray right
53 */
54 while(right < n && arr[(right-1)+1] - x < x - arr[left]){
55 right++;
56 left++;
57 }
58
59 // cout << left << ", " << right << endl;
60
61 return vector<int>(arr.begin()+left , arr.begin()+right);
62 }
63};
64
65//Approach 2: Binary Search and Two Pointers
66//similar as above, cleaner
67//Runtime: 116 ms, faster than 38.21% of C++ online submissions for Find K Closest Elements.
68//Memory Usage: 31.2 MB, less than 16.67% of C++ online submissions for Find K Closest Elements.
69//time: O(logN+K), space: O(K)
70class Solution {
71public:
72 vector<int> findClosestElements(vector<int>& arr, int k, int x) {
73 //find element >= x
74 auto it = lower_bound(arr.begin(), arr.end(), x);
75
76 if(it == arr.begin()){
77 return vector<int>(arr.begin(), arr.begin()+k);
78 }else if(it == arr.end()){
79 return vector<int>(arr.end() - k, arr.end());
80 }
81
82 int n = arr.size();
83 int idx = it - arr.begin();
84 int left, right;
85 /*
86 search range suggested by solution is [idx-(k-1), idx+(k-1)],
87 but idx will not necessary be included in final subarray,
88 so here expand the search range
89 */
90 left = max(0, idx - k); //max(0, idx - (k-1));
91 right = min(n-1, idx + k); //min(n-1, idx + (k-1));
92 //[left, right]: both inclusive
93
94 // cout << left << ", " << right << " ";
95
96 while(right - left + 1 > k){
97 if(x - arr[left] <= arr[right] - x){
98 right--;
99 }else if(arr[right] - x < x - arr[left]){
100 left++;
101 }
102 // cout << left << ", " << right << " ";
103 }
104
105 return vector<int>(arr.begin()+left , arr.begin()+right+1);
106 }
107};
108
109
110//Approach 1: sort
111//Runtime: 284 ms, faster than 7.52% of C++ online submissions for Find K Closest Elements.
112//Memory Usage: 31.2 MB, less than 16.67% of C++ online submissions for Find K Closest Elements.
113//time: O(NlogN), space: O(k)
114class Solution {
115public:
116 vector<int> findClosestElements(vector<int>& arr, int k, int x) {
117 sort(arr.begin(), arr.end(), [x](const int& a, const int& b){
118 return (abs(a-x) == abs(b-x)) ? a < b : (abs(a-x) < abs(b-x));
119 });
120
121 vector<int> ans(arr.begin(), arr.begin()+k);
122 //now ans is not ascending, need to sort it again
123
124 sort(ans.begin(), ans.end());
125
126 return ans;
127 }
128};Cost
Complexity
Time
O(n) to O(n log n), depending on the dominant loop or data structure operation
Dominated by the main traversal, recursion, or data-structure operations in the code.
Space
O(n) in the usual case for auxiliary containers or recursion
Auxiliary state plus the answer structure where the problem requires one.