配列が与えられた時、N個からなる最大値/最小値を求める

ABC 62-Dの3N Numbersを解いてるときに、N個からなる最大値/最小値を求める必要があったのでメモしときます。

Priority Queueを使う。

Priority Queueで実装できます。基本的なロジックは

  1. 最初のN個をPQにいれる。
  2. PQの中にある合計を記録する変数(sum)を用意する。
  3. N個目から先は
    1. i番目の数をPQにいれる。sumをアップデートする。
    2. PQから最小値を取り出す。取り出した値をsumから引く。

これを、配列の最後の要素にたどり着くまで繰り返す。新しい数字が入ってくるたびに、その集合の最小値が取り除かれるため、合計は常にN個からなる最大値になります。

コードはこんな感じです。

PriorityQueue<Integer> pq = new PriorityQueue<Integer>();
long sum = 0;
long max = Long.MIN_VALUE;

for (int i = 0; i < n; i++) {
    pq.add(nums[i]);
    sum += nums[i];
}

max = Math.max(max, sum);

for (int i = n; i < l; i++) {
    // Add an element
    pq.add(nums[i]);
    sum += nums[i];
    // Remove the minimum
    int min = pq.poll();
    sum -= min;
    // Update max
    max = Math.max(max, sum);
}

最小値

N個からなる最小値を求めたい場合、PriorityQueueの初期化のときに逆順を指定すればいいです。変数名などを変える必要がありますが、PQの初期化をこのようにすればN個からなる最小値を上のコードのまま求められます。PriorityQueue<Integer> pq = new PriorityQueue<Integer>(Collections.reverseOrder());