C++で奇数長のすべての部分リストの中央値の合計を求めるプログラム

問題の概要

数値のリスト nums が与えられたとき、そのリストから作れる「奇数の長さの部分リスト（サブリスト）」すべてを対象に、それぞれの中央値を求めて合計した値を計算するプログラムを考えてみましょう。

例えば、nums = [2, 4, 6, 3] が入力である場合を見てみます。奇数長の部分リストは [2]、[4]、[6]、[3]、[2, 4, 6]、[4, 6, 3] の6つです。それぞれの中央値は 2、4、6、3、4、4 となるため、合計は 2 + 4 + 6 + 3 + 4 + 4 = 23 になります。

解法の考え方：2つのヒープで中央値を管理

この問題を効率的に解くには、優先度付きキュー（ヒープ）を2つ組み合わせる手法が有効です。

• que_max（最大ヒープ）：数列の小さい側半分の要素を保持し、先頭に中央値が配置される
• que_min（最小ヒープ）：大きい側半分の要素を保持する

要素を追加するたびに、2つのヒープのサイズバランスと先頭同士の大小関係を調整すれば、常に先頭から中央値を即座に取得できます。

アルゴリズムの手順

1. 答えを格納する変数 ret を 0 で初期化します。
2. 部分リストの開始位置 i を 0 から順に末尾まで動かし、各ループで次の処理を行います。
   • 最大ヒープ que_max と最小ヒープ que_min を新しく用意します。
   • 終了位置 j を i から順に末尾まで動かしながら、以下を実行します。
     • nums[j] を que_max に挿入します。
     • que_max のサイズが que_min より 2 以上大きくなったら、que_max の先頭要素を que_min へ移動してバランスを取ります。
     • que_min が空ではなく、かつ que_max の先頭が que_min の先頭より大きい場合は、両者の先頭要素を取り出して入れ替えます。これにより「小さい側半分」と「大きい側半分」の境界が正しく維持されます。
     • i と j の偶奇が一致する場合（＝部分リストの長さが奇数の場合）、que_max の先頭（中央値）を ret に加算します。
3. すべての開始位置の処理が完了したら、ret を返します。

C++による実装例

#include <bits/stdc++.h>
using namespace std;
int solve(vector<int>& nums) {
    int ret = 0;
    for (int i = 0; i < nums.size(); i++) {
        priority_queue<int> que_max;
        priority_queue<int, vector<int>, greater<int>> que_min;
        for (int j = i; j < nums.size(); j++) {
            que_max.push(nums[j]);
            while (que_max.size() - que_min.size() >= 2) {
                que_min.push(que_max.top());
                que_max.pop();
            }
            while (que_min.size() && que_max.top() > que_min.top()) {
                int a = que_max.top();
                que_max.pop();
                int b = que_min.top();
                que_min.pop();
                que_max.push(b);
                que_min.push(a);
            }
            if (i % 2 == j % 2) {
                ret += que_max.top();
            }
        }
    }
    return ret;
}
int main(){
    vector<int> v = {2, 4, 6, 3};
    cout << solve(v);
}

入力

{2, 4, 6, 3}

出力

23

計算量とまとめ

外側の二重ループですべての開始位置と終了位置の組み合わせを走査し、各ステップでヒープ操作（O(log n)）を行うため、全体の計算量は O(n² log n) となります。ヒープを2つ併用することで、部分リストが伸びるたびに要素を並べ替え直すことなく中央値を追跡でき、部分リストごとにソートする素朴な方法よりも効率的です。また、奇数長の判定に「i と j の偶奇の一致」というシンプルな条件を利用している点も、この実装のポイントと言えるでしょう。