Given a sorted integer array, find the k closest elements to x in the array where k and x are given positive integers.
The integer x may or may not be present in the input array. If x is less than or equal to the first element in the input array, return first k elements. Similarly, if x is more than or equal to the last element in the input array, return the last k elements. The returned elements should be in the same order as present in the input array.
For example,
The idea is to do a linear search to find the insertion point i. The insertion point is defined as the point at which the key x would be inserted into the array, i.e., the index of the first element greater than the key, or the array’s size if all elements in the array are less than the specified key. Then, compare the elements around the insertion point i to get the first k closest elements. The time complexity of this solution is O(n).
We can also find the insertion point i using binary search algorithm which runs in O(log(n)) time. Since finding the k closest elements takes O(k) time, the overall time complexity of this solution is O(log(n) + k). Below is the implementation in C, C++, Java and Python based on the this idea.
C
|
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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
#include <stdio.h>
#include <stdlib.h>
// Function to search the specified array `arr` for the key `x`
// using the binary search algorithm
int binarySearch(int arr[], int n, int x)
{
int low = 0;
int high = n – 1;
while (low <= high)
{
int mid = low + (high – low) / 2;
if (arr[mid] < x) {
low = mid + 1;
}
else if (arr[mid] > x) {
high = mid – 1;
}
else {
return mid; // key found
}
}
return low; // key not found
}
// Function to find the `k` closest elements to `x` in a sorted integer array `arr`
void findKClosestElements(int arr[], int x, int k, int n)
{
// find the insertion point using binary search algorithm
int i = binarySearch(arr, n, x);
int left = i – 1;
int right = i;
// run `k` times
while (k— > 0)
{
// compare the elements on both sides of the insertion point `i`
// to get the first `k` closest elements
if (left < 0 || (right < n &&
abs(arr[left] – x) > abs(arr[right] – x))) {
right++;
}
else {
left—;
}
}
// print `k` closest elements
left++;
while (left < right)
{
printf(“%d “, arr[left]);
left++;
}
}
int main(void)
{
int arr[] = { 10, 12, 15, 17, 18, 20, 25 };
int n = sizeof(arr) / sizeof(arr[0]);
int x = 16, k = 4;
findKClosestElements(arr, x, k, n);
return 0;
}
|
Output:
12 15 17 18
C++
|
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
42
43
44
45
|
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
// Function to find the `k` closest elements to `x` in a sorted integer array `arr`
vector<int> findKClosestElements(vector<int> &arr, int x, int k)
{
// find the insertion point using binary search algorithm
int i = std::lower_bound(arr.begin(), arr.end(), x) – arr.begin();
int left = i – 1;
int right = i;
// run `k` times
while (k— > 0)
{
// compare the elements on both sides of the insertion point `i`
// to get the first `k` closest elements
if (left < 0 || (right < arr.size() &&
abs(arr[left] – x) > abs(arr[right] – x))) {
right++;
}
else {
left—;
}
}
// return `k` closest elements
return vector<int>(arr.begin() + left + 1, arr.begin() + right);
}
int main()
{
vector<int> arr = { 10, 12, 15, 17, 18, 20, 25 };
int x = 16, k = 4;
vector<int> result = findKClosestElements(arr, x, k);
for (int i: result) {
cout << i << ” “;
}
return 0;
}
|
Output:
12 15 17 18
Java
|
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
42
43
44
45
46
47
48
|
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
public class Main
{
// Function to find the `k` closest elements to `x` in a sorted list `arr`
public static List<Integer> findKClosestElements(List<Integer> arr, int k, int x)
{
// find the insertion point using binary search algorithm
int i = Collections.binarySearch(arr, x);
// Collections.binarySearch() returns `-(insertion point) – 1`
// if key is not contained in the list
if (i < 0) {
i = –(i + 1);
}
int left = i – 1;
int right = i;
// run `k` times
while (k— > 0)
{
// compare the elements on both sides of the insertion point `i`
// to get the first `k` closest elements
if (left < 0 || (right < arr.size() &&
Math.abs(arr.get(left) – x) > Math.abs(arr.get(right) – x))) {
right++;
}
else {
left—;
}
}
// return `k` closest elements
return arr.subList(left + 1, right);
}
public static void main(String[] args) {
List<Integer> arr = Arrays.asList(10, 12, 15, 17, 18, 20, 25);
int x = 16, k = 4;
System.out.println(findKClosestElements(arr, k, x));
}
}
|
Output:
[12, 15, 17, 18]
Python
|
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
42
43
44
45
46
47
48
49
50
|
# Function to search the specified array `arr` for the key `x`
# using the binary search algorithm
def binarySearch(arr, x):
low = 0
high = len(arr) – 1
while low <= high:
mid = low + (high – low) // 2
if arr[mid] < x:
low = mid + 1
elif arr[mid] > x:
high = mid – 1
else:
return mid # key found
return low # key not found
# Function to find the `k` closest elements to `x` in a sorted integer array `arr`
def findKClosestElements(arr, x, k):
# find the insertion point using binary search algorithm
i = binarySearch(arr, x)
left = i – 1
right = i
# run `k` times
while k > 0:
# compare the elements on both sides of the insertion point `i`
# to get the first `k` closest elements
if left < 0 or (right < len(arr) and abs(arr[left] – x) > abs(arr[right] – x)):
right = right + 1
else:
left = left – 1
k = k – 1
# return `k` closest elements
return arr[left+1: right]
if __name__ == ‘__main__’:
arr = [10, 12, 15, 17, 18, 20, 25]
x = 16
k = 4
print(findKClosestElements(arr, x, k))
|
Output:
[12, 15, 17, 18]
Download Run Code
Output:
12 15 17 18
Download Run Code
Output:
12 15 17 18
Download Run Code
Output:
[12, 15, 17, 18]
Download Run Code
Output:
[12, 15, 17, 18]
The above solution to do a binary search to find the insertion point and then tries to find the k closest elements. However, we can combine the whole logic in a single binary search routine. Here’s how the algorithm would look like in C, Java, and Python.
C
|
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
|
#include <stdio.h>
#include <stdlib.h>
// Function to find the `k` closest elements to `x` in a sorted integer array `arr`
void findKClosestElements(int arr[], int x, int k, int n)
{
int left = 0;
int right = n;
while (right – left >= k)
{
if (abs(arr[left] – x) > abs(arr[right] – x)) {
left++;
}
else {
right—;
}
}
// print `k` closest elements
while (left <= right)
{
printf(“%d “, arr[left]);
left++;
}
}
int main(void)
{
int arr[] = { 10, 12, 15, 17, 18, 20, 25 };
int n = sizeof(arr) / sizeof(arr[0]);
int x = 16, k = 4;
findKClosestElements(arr, x, k, n);
return 0;
}
|
Output:
12 15 17 18
Java
|
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
|
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;
public class Main
{
// Function to find the `k` closest elements to `x` in a sorted integer array `arr`
public static List<Integer> findKClosestElements(int[] arr, int k, int x)
{
int left = 0;
int right = arr.length – 1;
while (right – left >= k)
{
if (Math.abs(arr[left] – x) > Math.abs(arr[right] – x)) {
left++;
} else {
right—;
}
}
return Arrays.stream(arr, left, right + 1).boxed()
.collect(Collectors.toList());
}
public static void main(String[] args) {
int[] arr = {10, 12, 15, 17, 18, 20, 25 };
int x = 16, k = 4;
System.out.println(findKClosestElements(arr, k, x));
}
}
|
Output:
[12, 15, 17, 18]
Python
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
# Function to find the `k` closest elements to `x` in a sorted integer array `arr`
def findKClosestElements(arr, k, x):
left = 0
right = len(arr) – 1
while right – left >= k:
if abs(arr[left] – x) > abs(arr[right] – x):
left = left + 1
else:
right = right – 1
return arr[left:left + k]
if __name__ == ‘__main__’:
arr = [10, 12, 15, 17, 18, 20, 25]
x = 16
k = 4
print(findKClosestElements(arr, k, x))
|
Output:
[12, 15, 17, 18]
Download Run Code
Output:
12 15 17 18
Download Run Code
Output:
[12, 15, 17, 18]
Download Run Code
Output:
[12, 15, 17, 18]
The time complexity of above solution is O(n) and auxiliary space used by the program is O(1).