658.find-k-closest-elements.md

問題

https://leetcode.com/problems/find-k-closest-elements/

回答

初回の回答:

まず二分探索で x と等しい要素の index を見つけ、そこから左右に k 個分抜き出す方針で実装した。この実装ではパスしなかった。「x と最も近い要素を k 個」という条件なので、必ずしも x と等しい要素がなくても OK だということが考慮されていなかった。また、同じ値の要素が複数個ある（$[1,1,1,10,10,10]$など）場合も考慮されていなかった。

impl Solution {
    pub fn find_closest_elements(arr: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
        fn binary_search(arr: &Vec<i32>, x: &i32) -> i32 {
            let mut l = 0;
            let mut r = arr.len();
            while l < r {
                let m = l + (r - l) / 2;
                match arr[m].cmp(&x) {
                    std::cmp::Ordering::Equal => return m as i32,
                    std::cmp::Ordering::Greater => r = m,
                    std::cmp::Ordering::Less => l = m + 1,
                }
            }
            -1
        }

        if arr.last().unwrap() < &x {
            return arr[(arr.len() - k as usize)..arr.len()].to_vec()
        } else if arr.first().unwrap() > &x {
            return arr[0..k as usize].to_vec()
        }

        let i = binary_search(&arr, &x);
        let d = (k - 1) / 2;
        let r = (k - 1) % 2;

        let mut start = i - (d + r);
        let mut end = i + d;
        let last_index = (arr.len() - 1) as i32;

        if start < 0 {
            end += start * -1;
            start = 0;
        } else if end > last_index {
            start -= (end - last_index);
            end = last_index;
        }

        arr[start as usize..=end as usize].to_vec()
    }
}

二回目の回答:

二分探索を複数回行って、その都度最も近い要素を 1 個ないし複数個抜き出す方針で実装した。これはパスした。計算量的にも二分探索を複数回行うので多分 $N\mathcal{O}(\log{_2}n)$ 。

impl Solution {
    pub fn find_closest_elements(arr: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
        //  return the first index of the most closest or the exact integer
        fn find_closest(arr: &mut Vec<i32>, x: &i32) -> Vec<i32> {
            let mut l = 0;
            let mut r = arr.len() - 1;
            while l < r {
                let m = l + (r - l) / 2;
                if arr[m] >= *x {
                    r = m;
                } else {
                    l = m + 1;
                }
            }

            let start = if arr[l] == *x || arr.get(l - 1).is_none() {
                l
            } else {
                let lower_diff = (arr[l - 1] - *x).abs();
                let upper_diff = (arr[l] - *x).abs();
                match lower_diff.cmp(&upper_diff) {
                    std::cmp::Ordering::Equal => if arr[l - 1] < arr[l] { l - 1 } else { l },
                    std::cmp::Ordering::Greater => l,
                    std::cmp::Ordering::Less => l - 1
                }
            };

            r = arr.len();
            while l < r {
                let m = l + (r - l) / 2;
                match arr[m].cmp(&arr[start]) {
                    std::cmp::Ordering::Equal => l = m + 1,
                    _ => r = m
                }
            }
            let end = l - 1;
            arr.drain(start..=end).collect()
        }

        let mut arr = arr;
        let mut v = Vec::new();
        while v.len() < k as usize {
            v.append(&mut find_closest(&mut arr, &x))
        }
        v.sort();
        v[0..k as usize].to_vec()
    }
}

三回目の回答:

k 個抜き出す分の範囲の始端を求めるという考え方。配列が昇順にソートされていることを利用すれば、たしかに始端だけ分かれば問題は解けるしシンプルだけど、どうやって求めるかを考えるのは難しい。

$arr[m]$ が仮に始端だとすると、取り出す範囲は $m..m + k - 1$ で $m+k$ は範囲外であって、これは逆も然りである。したがって、 $arr[m]$ か $arr[m + k]$ どちらかは必ず含まれないことはたしか。なので、$arr[m]$ か $arr[m + k]$ どちらかが $x$ により近いかを判定して、範囲を徐々に絞っていく。

$\mathcal{O}(\log_2n)$ で探索するためにステップごとに範囲を半分に絞っていく。以下に $x, arr[m], arr[m+k]$ の大小関係、距離のパターンを表すと、左右どちらに絞るべきかがわかってくる。

• $[---x---arr[m]---arr[m + k]----]$
  → $r = m$ にして左半分に絞る。そうしないと $x$ により近い要素を含められない。

• $[---arr[m]-x-----arr[m + k]----]$
  → $r = m$ にして左半分に絞る。$arr[m]$ より多分少し左側に始端がある。

• $[---arr[m]-----x-arr[m + k]----]$
  → $l = m + 1$ にして右半分に絞る。$arr[m]$ より多分少し右側に始端がある。

• $[---arr[m]---arr[m + k]---x----]$
  → $l = m + 1$ にして右半分に絞る。明らかに $arr[m]$ より右に始端がある。

impl Solution {
    pub fn find_closest_elements(arr: Vec<i32>, k: i32, x: i32) -> Vec<i32> {
        // find a left bound and take k elements from there
        let k = k as usize;
        let mut l = 0;
        let mut r = arr.len() - k;
        while l < r {
            let m = l + (r - l) / 2;
            if x - arr[m] > arr[m + k] - x {
                l = m + 1;
            } else {
                r = m;
            }
        }
        arr[l..(l + k)].to_vec()
   }
}