LeetCode 658. Find K Closest Elements

LeetCode 658. Find K Closest Elements

這一題我個人認為難度相對高一點,而且答案不是很好想,實作的時候更要注意邊界,不過整體來說是蠻有技巧的一題,如果對於 Heap 有想要更深一層的認識,我覺得是可以挑戰看看。

先來看一下這個問題原文長怎樣?


Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

  • |a - x| < |b - x|, or
  • |a - x| == |b - x| and a < b

Example 1:

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

Example 2:

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

Constraints:

  • 1 <= k <= arr.length
  • 1<=arr.length<=1041 <= arr.length <= 10^4
  • arr is sorted in ascending order.
  • 104<=arr[i],x<=104-10^4 <= arr[i], x <= 10^4

首先一開始我們解釋一下題目,給你一個排序過的陣列並且給你兩個值 k 與 x,請找出陣列中離 x 最近的 k 個元素,回傳陣列必須是升冪陣列。

一樣是丟到 Heap 內讓它幫我們完成排序後,我們這邊用一個小技巧就是將 Heap 設定 K 個元素,接著透過二元排序法將其找到 x 的位置,如果 x 不在陣列內,則會找到最接近他的位置,這樣一來就可以把範圍匡出來。

完整的解答。

class Solution {
    fun findClosestElements(arr: IntArray, k: Int, x: Int): List<Int> {
        val foundIndex = binearSearch(arr, x)
        //target is X ==> [K elements] X
        var leftRange = foundIndex - k
        //target is X ==> X [K elements]
        var rightRange = foundIndex + k
        //check left and right range is in array?
        leftRange = Math.max(leftRange, 0)
        rightRange = Math.min(rightRange, arr.size-1)
        
        //key, value => index, distance
        val minHeap = PriorityQueue<Pair<Int,Int>>(){p1,p2->
            p1.first - p2.first
        }
        for(i in leftRange..rightRange){
            minHeap.add(Pair(Math.abs(arr[i]-x), i))
        }
        val result = arrayListOf<Int>()
        for(i in 0..k){
            //get index
            val index = minHeap.poll()?.first?:0
            result.add(arr[index])
        }
        return result
    }
    
    //return left index
    fun binearSearch(nums:IntArray, x:Int):Int{
        var right = 0
        var left = nums.size - 1
        while(right <= left){
            var mid = (right + left)/2
            if(nums[mid] == x){
                return mid
            } else if(nums[mid] < x){
                right = mid-1
            } else{
                left = mid+1
            }
        }
        return left
    }
}