[Google Interview] Find the k Closest Numbers in a Sorted Array

[toc]

? This is one of the Google interview questions and reported by programmers across the globe as one of the commonly asked questions during interviews. So, can you give the optimal solution to this problem?

Problem Formulation

Given an integer array or Python list nums, an integer value x and k.

Find and return the k closest numbers to the input x in the array.

⚠️ Constraints: You can assume that k is a number between 1 and the length of the nums list.

• 1 <= k <= nums.length
  • Therefore, it is implicitly ensured that the list nums has at least one element and there must always be exactly one solution.
• nums is sorted in ascending order.

?Examples

Let’s have a look at some examples to improve our understanding of this problem.

Example 1

Input:  [8, 10, 12, 15, 18, 20], k = 4, x = 15
Output: [10, 12, 15, 18]


Example 2

Input:  [4, 6, 8, 9], k = 3, x = 7
Output: [6, 8, 9]


Example 3

Input:  [2, 3, 5, 6, 7], k = 1, x = 4
Output: [3]


Example 4

Input:  [5], k = 1, x = 4
Output: [5]


Example 5

Input:  [10, 11, 12, 13, 15, 16], k = 1, x = 15
Output: [15]

?️An Easy Approach: Binary Search

The most straightforward solution to this problem is to use binary search as follows:

• First, use the binary search algorithm to find the insertion point (The insertion point is the point where the integer x can be placed or inserted within the sorted list). The elements before this point are smaller, while the elements after it are greater.
• Then, compare the elements around this point to find the k closest numbers.

Let’s have a look at the code that accomplishes that:

def binary(nums, x):
    start = 0
    end = len(nums) - 1

    while start <= end:
        mid = start + ((end - start) // 2)
        if nums[mid] < x:
            start = mid + 1
        elif nums[mid] > x:
            end = mid - 1
        else:
            return mid
    return start


def k_close(nums, x, k):
    no = binary(nums, x)
    lhs = no - 1
    rhs = no
    while k > 0:
        if lhs < 0 or (rhs < len(nums) and abs(nums[lhs] - x) > abs(nums[rhs] - x)):
            rhs = rhs + 1
        else:
            lhs = lhs - 1
        k = k - 1
    return nums[lhs + 1: rhs]

Let’s run this code on our examples:

# Example 1
nums = [8, 10, 12, 15, 18, 20]
k = 4
x = 15
print(k_close(nums, x, k))
# [10, 12, 15, 18]

# Example 2
nums = [4, 6, 8, 9]
k = 3
x = 7
print(k_close(nums, x, k))
# [6, 8, 9]

# Example 3
nums = [2, 3, 5, 6, 7]
k = 1
x = 4
print(k_close(nums, x, k))
# [3]

# Example 4
nums = [5]
k = 1
x = 5
print(k_close(nums, x, k))
# [5]

# Example 5
nums = [10, 11, 12, 13, 15, 16]
k = 1
x = 15
print(k_close(nums, x, k))
# [15]

Hurrah! ? The code passed all the test cases.

❖ Analysis: The code consists of two functions: binary search and finding the k closest number.  The binary search algorithm has the time complexity of O(log(n)). The time complexity for finding the k closest numbers is O(k). Hence, the total complexity of this code becomes O(log n + k).

? Tidbit: The double-backslash // operator performs integer division and the single-backslash / operator performs float division. An example for integer division is 40//11 = 3. An example for float division is 40/11 = 3.6363636363636362.

❖ Discussion: We have performed a lot of extra work in the above approach as we performed the binary search for the whole list inside one method and then we used another method for computing the k closest numbers to the given value x. Is there a better way to deal with this problem?

?️The Optimal Solution

The better way would be to combine both the methods and generate an optimal solution. The main idea of this algorithm is to find out the lower bound for the given k length range. The numbers between “left” and “right” are the candidates of the lower bound.

❖ Approach: Assuming that A[mid] ~ A[mid + k] represents a sliding window, we compare the distance between x - A[mid] and A[mid + k] - x. Now let’s consider the following cases:

• as long as x - A[mid] > A[mid + k] - x, we need to move window go right.
• otherwise, we need to move the window to the left.

Here’s an example that illustrates the algorithm:

Now, let us have a look at the code:

def k_close(nums, x, k):
    left, right = 0, len(nums) - k
    while left < right:
        mid = (left + right) // 2
        if x - nums[mid] > nums[mid + k] - x:
            left = mid + 1
        else:
            right = mid
    return nums[left:left + k]

❖ Discussion:

• The if condition x - A[mid] > A[mid + k] - x is used to compare A[mid] and A[mid+k] to check which is closer to x.
• If A[mid] is closer to x, then A[mid+k] can never be in the k length range. So you can definitely remove all (A[mid+1], A[mid+2], A[mid+3]…) from the candidate list by setting right=mid.
• If A[mid+k] is closer to x, then A[mid] can never be in the k length range. So you can remove all (….A[mid-2], A[mid-1], A[mid]) from the candidates list by setting left=mid+1.
• Once you are left with only one candidate, i.e., left==right, you got our final lower bound and now you can return the k closest numbers by slicing the list.

❖ Test Cases:

❖ Time Complexity Analysis:

The operations to shift the pointers and compute the closest numbers within the loop have a time complexity of O(log (n-k)) and the time complexity to slice the list and return the desired output is O(k). Thus, the total time complexity of this algorithm is O(log(n-k)+k).

Let’s consider the following example to analyze the time complexity:

Given:

nums = [10, 11, 12, 13, 15, 16, 18, 19, 20, 22, 23]
k = 3
x = 15

• Let’s assume that the length of nums is ‘n’.  Thus, we shrink the pointers/window by (n-k) steps as shown in the table below. Thus, the while loop has a complexity of O(log(n – k)).
  • In the above example, n = 11 and k = 3. Thus, the while loop undergoes log(n-k) iterations, i.e., log(11- 3) ⇒ log 8 = 3 iterations.

• Finally, when we return the sliced list which represents the window containing the k closest neighbors, it takes O(k) time.
•  Hence, the overall complexity becomes O(log(n – k) + k).

?️A Bonus Solution: Using bisect and two pointers

Before discussing this approach, you need to understand what the bisect.bisect_left does. In a coding interview, you can usually assume you have access to basic external functionality. Here’s a basic recap on the idea of the bisect method:

? Bisect Recap:
◆ The purpose of the Bisect algorithms is to find the index/position of a required element within a given list where the element has to be inserted within the list. Therefore, it helps to keep the list sorted after the insertion is complete.
◆ bisect_left method of the bisect module is used to find the index of the target element in the sorted list. If the element is already present in the list then the leftmost position where the element can be inserted within the list is returned.

❖ Approach: The basic idea of this solution is to locate the insertion point for value x by using the bisect.bisect_left function in the module. Then, we will use two pointers to find the k closest elements. 

Let’s have a look at the code:

import bisect


def k_close(nums, x, k):
    pos = bisect.bisect_left(nums, x)
    left, right = pos - 1, pos
    while k:
        if right >= len(nums) or \
                (left >= 0 and abs(nums[left] - x) <= abs(nums[right] - x)):
            left -= 1
        else:
            right += 1
        k -= 1
    return nums[left + 1:right]

❖ Runtime analysis:
The Bisect function works by repeatedly halving a list. This means it has a running time of O(log n). The algorithm takes O(k) time to search the k closest numbers. Hence, the total complexity for this solution is O(log n+ k).

Note: This is a fantastic approach that one could come up with during an interview. However, it is to be noted that interviewers might actually ask you to implement an algorithm that uses binary search.

Conclusion

I hope you enjoyed this coding interview question. Please stay tuned and subscribe for more interesting coding problems.

Post Credits: Shubham Sayon and Rashi Agarwal

Recommended:  Finxter Computer Science Academy

• Do you want to master the most popular Python IDE fast?
• This course will take you from beginner to expert in PyCharm in ~90 minutes.
• For any software developer, it is crucial to master the IDE well, to write, test and debug high-quality code with little effort.

Join the PyCharm Masterclass now, and master PyCharm by tomorrow!