BFE.devでのAI解答例 - JavaScript コーディング問題
45. ソートしてない配列にてK番目に大きい要素を見つける

解法1: Quick Select

Quick Selectは、Quick Sortの応用であり、特定の要素をソートする必要がない場合に最適です。これは、一定の時間内にK番目大きい数字を見つけることができることが知られています。

アルゴリズムのアイデアは以下の通りです。

• 配列の中からピボットを選んで、それより大きい数字を右に、小さい数字を左に分割します。
• 分割後、ピボットがソートされた配列の何番目に位置するかを調べます。
• もし、ピボットがK番目の数字に位置しているならそれを返します。
• そうでなければ、右側にあるか左側にあるかを判断して再帰的に行います。

以下が実装例です。

function quickSelect(arr, k) {
  const pivotIndex = Math.floor(Math.random() * arr.length);
  const pivot = arr[pivotIndex];
  const left = [];
  const right = [];

  for (let i = 0; i < arr.length; i++) {
    if (i !== pivotIndex) {
      arr[i] < pivot ? left.push(arr[i]) : right.push(arr[i]);
    }
  }

  if (right.length >= k) {
    return quickSelect(right, k);
  } else if (right.length === k - 1) {
    return pivot;
  } else {
    return quickSelect(left, k - right.length - 1);
  }
}

このアルゴリズムでは、ピボットをランダムに選ぶことで最悪の場合でもO(n)の時間でK番目の数字を見つけることができます。

解法2: Heap

Heapを使ってもK番目の数字を見つけることができます。Heapは、最大/最小の要素をO(1)で見つけることができるという特徴がありますが、他の要素を見つけるのには線形時間がかかるため、この問題には不向きです。

しかし、K番目の数字を見つけるためにはK個の最大/最小の数字に興味がありますので、Kサイズの最小Heap（最大Heap）を使って部分的に配列をソートすることができます。

以下が実装例です。

function findKthLargest(arr, k) {
  const heap = new MinHeap();

  for (let i = 0; i < arr.length; i++) {
    heap.insert(arr[i]);
    if (heap.size() > k) {
      heap.extract();
    }
  }

  return heap.extract();
}

class MinHeap {
  constructor() {
    this.heap = [null];
  }

  size() {
    return this.heap.length - 1;
  }

  insert(value) {
    this.heap.push(value);
    this.bubbleUp(this.size());
  }

  extract() {
    if (this.size() < 1) {
      return null;
    }

    const value = this.heap[1];
    const lastValue = this.heap.pop();
    if (this.size() > 0) {
      this.heap[1] = lastValue;
      this.bubbleDown(1);
    }
    return value;
  }

  bubbleUp(index) {
    const value = this.heap[index];

    while (index > 1 && value < this.heap[Math.floor(index / 2)]) {
      this.heap[index] = this.heap[Math.floor(index / 2)];
      index = Math.floor(index / 2);
    }

    this.heap[index] = value;
  }

  bubbleDown(index) {
    const value = this.heap[index];

    while (2 * index <= this.size()) {
      let childIndex = 2 * index;
      if (childIndex + 1 <= this.size() && this.heap[childIndex + 1] < this.heap[childIndex]) {
        childIndex++;
      }
      if (value <= this.heap[childIndex]) {
        break;
      }

      this.heap[index] = this.heap[childIndex];
      index = childIndex;
    }

    this.heap[index] = value;
  }
}

このアルゴリズムは、配列の長さがNのとき、最初にK個の数字をHeapに挿入するためにO(KlogK)の時間がかかります。その後、残りの(N - K)個の数字を一つずつ取り出してHeapに挿入します。Heapの最小値には常にK番目大きい数字が含まれるため、配列からの取り出しは最大でもN - K回行います。O(logK)時間で最小値を取り出せるため、結局のところこのアルゴリズムの計算量はO(KlogK + (N - K)logK) = O(NlogK)となります。