サブリストを1つ削除して、k未満とk超の要素数を等しくする最長リストを求めるC++プログラム

数列 nums ともう一つの整数 k が与えられているとします。ここで、リストから連続する部分リスト（サブリスト）を高々一度だけ削除できるものとします。削除後のリストの中で、「k より厳密に小さい数」と「k より厳密に大きい数」の個数が等しくなるような、最長のリストの長さを求めるのが目的です。

たとえば、入力が nums = [6, 10, 8, 9, 3, 5]、k = 6 の場合を考えてみましょう。このとき出力は 5 になります。部分リスト [9] を削除すると [6, 10, 8, 3, 5] が得られ、6 より小さい数は [3, 5] の2つ、6 より大きい数は [10, 8] の2つとなり、個数が一致するためです。

解法のアプローチ

この問題は「累積バランス」とハッシュマップを組み合わせることで効率的に解けます。考え方のポイントは次の通りです。

• バランス値の計算： リストを左から走査し、k 未満の数が出たら +1、k 超の数が出たら −1 として累積値を配列 v に記録します。v の各要素は「その位置までのバランス」を表します。
• 全体バランスの確認： v の最後の要素（全体のバランス）が 0 であれば削除は不要なので、リスト全体の長さをそのまま返します。
• 削除すべき区間の特定： 全体バランスが delta（≠ 0）の場合、「v[j] − v[i] = delta」となる区間 (i, j] を削除すればバランスが 0 になります。つまり、差が delta となる2つのインデックスを見つけ、その間隔が最小となる区間を探します。
• マップによる高速化： 各バランス値が最初に出現した位置をマップ m に記録しながら一回の走査で処理することで、最小の削除区間を効率よく求められます。

具体的な手順は以下の通りです。

• nums と同じサイズ＋1 の配列 v を定義し、すべて 0 で初期化する
• cnt := 0 とする
• i := 0 から nums のサイズ未満の間、i を1ずつ増やしながら以下を実行する
  • nums[i] < k の場合：cnt を1増やす
  • そうでなく nums[i] > k の場合：cnt を1減らす
  • v[i + 1] = cnt とする
• v の最後の要素が 0 なら、nums のサイズを返す
• delta := v の最後の要素とする
• マップ m を定義する
• ans := 無限大とする
• i := 1 から v のサイズ以下の間、i を1ずつ増やしながら以下を実行する
  • m[v[i] − v の最後の要素] が 0 でない、または v[i] − v の最後の要素が 0 に等しい場合：
    • ans := ans と i − m[v[i] − v の最後の要素] の最小値
  • m[v[i]] := i とする
• ans が無限大のままなら 0 を返す
• そうでなければ nums のサイズ − ans を返す

それでは、実際の実装を見て理解を深めましょう。

C++実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int solve(vector<int>& nums, int k) {
        vector<int> v(nums.size() + 1, 0);
        int cnt = 0;
        for (int i = 0; i < nums.size(); ++i) {
            if (nums[i] < k)
                ++cnt;
            else if (nums[i] > k)
                --cnt;
            v[i + 1] = cnt;
        }
        if (v.back() == 0) return int(nums.size());
        int delta = v.back();
        map<int, int> m;
        int ans = INT_MAX;
        for (int i = 1; i <= v.size(); ++i) {
            if (m[v[i] - v.back()] != 0 || v[i] - v.back() == 0) {
                ans = min(ans, i - m[v[i] - v.back()]);
            }
            m[v[i]] = i;
        }
        if (ans == INT_MAX)
            return 0;
        else
            return int(nums.size() - ans);
    }
};
main(){
    Solution ob;
    vector<int> v = {6, 10, 8, 9, 3, 5};
    int k = 6;
    cout << ob.solve(v, k); }

入力

{6, 10, 8, 9, 3, 5}, 6

出力

5

まとめ

このアルゴリズムでは、まず累積バランスを前計算し、削除によって打ち消すべき差分 delta に着目します。同じバランス値同士の距離が削除候補区間の長さに対応するため、マップで最初の出現位置を記録しながら一回の走査だけで最小区間を特定できます。計算量は O(n log n)、空間計算量は O(n) であり、削除候補をすべて総当たりする O(n²) の素朴な手法よりも大幅に効率的です。