Question

658. Find K Closest Elements

中文文档

Description

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 <= 104
• arr is sorted in ascending order.
• -104 <= arr[i], x <= 104

Solutions

Solution 1: Sort

class Solution:
  def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
    arr.sort(key=lambda v: abs(v - x))
    return sorted(arr[:k])

class Solution {
  public List<Integer> findClosestElements(int[] arr, int k, int x) {
    List<Integer> ans = Arrays.stream(arr)
                .boxed()
                .sorted((a, b) -> {
                  int v = Math.abs(a - x) - Math.abs(b - x);
                  return v == 0 ? a - b : v;
                })
                .collect(Collectors.toList());
    ans = ans.subList(0, k);
    Collections.sort(ans);
    return ans;
  }
}

int target;

class Solution {
public:
  static bool cmp(int& a, int& b) {
    int v = abs(a - target) - abs(b - target);
    return v == 0 ? a < b : v < 0;
  }

  vector<int> findClosestElements(vector<int>& arr, int k, int x) {
    target = x;
    sort(arr.begin(), arr.end(), cmp);
    vector<int> ans(arr.begin(), arr.begin() + k);
    sort(ans.begin(), ans.end());
    return ans;
  }
};

func findClosestElements(arr []int, k int, x int) []int {
  sort.Slice(arr, func(i, j int) bool {
    v := abs(arr[i]-x) - abs(arr[j]-x)
    if v == 0 {
      return arr[i] < arr[j]
    }
    return v < 0
  })
  ans := arr[:k]
  sort.Ints(ans)
  return ans
}

func abs(x int) int {
  if x >= 0 {
    return x
  }
  return -x
}

function findClosestElements(arr: number[], k: number, x: number): number[] {
  let l = 0;
  let r = arr.length;
  while (r - l > k) {
    if (x - arr[l] <= arr[r - 1] - x) {
      --r;
    } else {
      ++l;
    }
  }
  return arr.slice(l, r);
}

impl Solution {
  pub fn find_closest_elements(arr: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
    let n = arr.len();
    let mut l = 0;
    let mut r = n;
    while r - l != (k as usize) {
      if x - arr[l] <= arr[r - 1] - x {
        r -= 1;
      } else {
        l += 1;
      }
    }
    arr[l..r].to_vec()
  }
}

Solution 2: Binary search

class Solution:
  def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
    l, r = 0, len(arr)
    while r - l > k:
      if x - arr[l] <= arr[r - 1] - x:
        r -= 1
      else:
        l += 1
    return arr[l:r]

class Solution {
  public List<Integer> findClosestElements(int[] arr, int k, int x) {
    int l = 0, r = arr.length;
    while (r - l > k) {
      if (x - arr[l] <= arr[r - 1] - x) {
        --r;
      } else {
        ++l;
      }
    }
    List<Integer> ans = new ArrayList<>();
    for (int i = l; i < r; ++i) {
      ans.add(arr[i]);
    }
    return ans;
  }
}

class Solution {
public:
  vector<int> findClosestElements(vector<int>& arr, int k, int x) {
    int l = 0, r = arr.size();
    while (r - l > k) {
      if (x - arr[l] <= arr[r - 1] - x) {
        --r;
      } else {
        ++l;
      }
    }
    return vector<int>(arr.begin() + l, arr.begin() + r);
  }
};

func findClosestElements(arr []int, k int, x int) []int {
  l, r := 0, len(arr)
  for r-l > k {
    if x-arr[l] <= arr[r-1]-x {
      r--
    } else {
      l++
    }
  }
  return arr[l:r]
}

function findClosestElements(arr: number[], k: number, x: number): number[] {
  let left = 0;
  let right = arr.length - k;
  while (left < right) {
    const mid = (left + right) >> 1;
    if (x - arr[mid] <= arr[mid + k] - x) {
      right = mid;
    } else {
      left = mid + 1;
    }
  }
  return arr.slice(left, left + k);
}

impl Solution {
  pub fn find_closest_elements(arr: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
    let k = k as usize;
    let n = arr.len();
    let mut left = 0;
    let mut right = n - k;
    while left < right {
      let mid = left + (right - left) / 2;
      if x - arr[mid] > arr[mid + k] - x {
        left = mid + 1;
      } else {
        right = mid;
      }
    }
    arr[left..left + k].to_vec()
  }
}

Solution 3

class Solution:
  def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
    left, right = 0, len(arr) - k
    while left < right:
      mid = (left + right) >> 1
      if x - arr[mid] <= arr[mid + k] - x:
        right = mid
      else:
        left = mid + 1
    return arr[left : left + k]

class Solution {
  public List<Integer> findClosestElements(int[] arr, int k, int x) {
    int left = 0;
    int right = arr.length - k;
    while (left < right) {
      int mid = (left + right) >> 1;
      if (x - arr[mid] <= arr[mid + k] - x) {
        right = mid;
      } else {
        left = mid + 1;
      }
    }
    List<Integer> ans = new ArrayList<>();
    for (int i = left; i < left + k; ++i) {
      ans.add(arr[i]);
    }
    return ans;
  }
}

class Solution {
public:
  vector<int> findClosestElements(vector<int>& arr, int k, int x) {
    int left = 0, right = arr.size() - k;
    while (left < right) {
      int mid = (left + right) >> 1;
      if (x - arr[mid] <= arr[mid + k] - x)
        right = mid;
      else
        left = mid + 1;
    }
    return vector<int>(arr.begin() + left, arr.begin() + left + k);
  }
};

func findClosestElements(arr []int, k int, x int) []int {
  left, right := 0, len(arr)-k
  for left < right {
    mid := (left + right) >> 1
    if x-arr[mid] <= arr[mid+k]-x {
      right = mid
    } else {
      left = mid + 1
    }
  }
  return arr[left : left+k]
}