Formatted question description: https://leetcode.ca/all/658.html

658. Find K Closest Elements

Level

Medium

Description

Given a sorted array, two integers k and x, find the k closest elements to x in the array. The result should also be sorted in ascending order. If there is a tie, the smaller elements are always preferred.

Example 1:

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

Output: [1,2,3,4]

Example 2:

Input: [1,2,3,4,5], k=4, x=-1

Output: [1,2,3,4]

Note:

  1. The value k is positive and will always be smaller than the length of the sorted array.
  2. Length of the given array is positive and will not exceed 104.
  3. Absolute value of elements in the array and x will not exceed 104.

Solution

Use binary search to find the index of x in the array, or the index at which x should be inserted in the array. Then use two pointers index1 and index2 that point to numbers less than x and greater than x respectively to determine which numbers are closest to x. If both numbers have the same absolute difference to x or the number at index1 has a smaller absolute difference, then add the number at index1 to the result list and move index1 to the left by 1 step. Otherwise, add the number at index2 to the result list and move index2 to the right by 1 step. If one pointer is out of the bound of the array, then only consider the other pointer. Repeat the process until k numbers are added to the result list. Finally, sort the result list and return.

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        List<Integer> closestElements = new ArrayList<Integer>();
        int count = 0;
        int length = arr.length;
        int index = binarySearch(arr, x);
        int index1 = -1, index2 = -1;
        if (index >= 0) {
            closestElements.add(arr[index]);
            count++;
            index1 = index - 1;
            index2 = index + 1;
        } else {
            index = -index - 1;
            index1 = index - 1;
            index2 = index;
        }
        while (count < k && index1 >= 0 && index2 < length) {
            int num1 = arr[index1], num2 = arr[index2];
            int compare = (num2 - x) - (x - num1);
            if (compare >= 0) {
                closestElements.add(num1);
                index1--;
            } else {
                closestElements.add(num2);
                index2++;
            }
            count++;
        }
        while (count < k && index1 >= 0) {
            int num1 = arr[index1];
            closestElements.add(num1);
            index1--;
            count++;
        }
        while (count < k && index2 < length) {
            int num2 = arr[index2];
            closestElements.add(num2);
            index2++;
            count++;
        }
        Collections.sort(closestElements);
        return closestElements;
    }

    public int binarySearch(int[] arr, int key) {
        int low = 0, high = arr.length - 1;
        while (low < high) {
            int mid = (high - low) / 2 + low;
            int num = arr[mid];
            if (num == key)
                return mid;
            else if (num > key)
                high = mid - 1;
            else
                low = mid + 1;
        }
        return -low - 1;     
    }
}

All Problems

All Solutions