C++で制約付き部分列の最大和を効率よく求める方法

問題の概要

整数型の配列 nums と整数 k が与えられたとき、次の条件を満たす空でない部分列の合計の最大値を求めます。すなわち、部分列の中で隣り合う任意の2つの数 nums[i] と nums[j](i < j)について、常に j - i <= k が成り立つことです。

ここでいう部分列とは、元の配列からいくつかの要素を削除し、残った要素を元の順序のまま保ったものを指します。

入力例

入力が [10, 2, -9, 5, 19]、k = 2 の場合を考えてみましょう。このとき最適な部分列は [10, 2, 5, 19] となり、その合計は 36 になります。

解き方のポイント

この問題は、動的計画法(DP)と「単調デック」によるスライディングウィンドウ最大値のテクニックを組み合わせることで、O(n) 時間で解けます。dp[i] を「インデックス i を必ず含む部分列の合計の最大値」と定義すると、遷移式は次のようになります。

dp[i] = nums[i] + max(0, dp[i-1], dp[i-2], ..., dp[i-k])

つまり、直近 k 個以内にある dp の値のうち最大のものを現在の要素に足し合わせます。ただし、直前までの累積値が負になっている場合は、そこで部分列を切り離して新しく始めたほうがよいため、max に 0 を含めておきます。この「範囲内の最大値」を毎回線形探索すると O(nk) かかってしまいますが、単調デックを使えばならし O(1) で取得できます。

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

• ret := -inf(答えを格納する変数)

• 配列 dp を定義し、与えられた配列の内容をそのままコピーする

• デック dq を1つ定義する

• v[0] を dq の先頭に挿入する

• n := v のサイズ、ret := v[0]

• i := 1 から開始し、i < n の間、i を1ずつ増やしながら以下を繰り返す:

  • i > k かつ dq の先頭要素が dp[i - k - 1] と等しい場合は、dq の先頭要素を削除する(ウィンドウの範囲外になった古い値を取り除く)

  • dp[i] := dp[i] と(dq が空なら dp[i]、空でなければ dp[i] + dq の先頭要素)の大きい方

  • dq が空でなく、dq の末尾要素が dp[i] より小さい間、dq の末尾要素を削除する(デック内の単調性を維持)

  • dp[i] を dq の末尾に挿入する

  • ret := ret と dp[i] の大きい方

• ret を返す

C++での実装例

理解を深めるために、実際のコードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
const int inf = 1e9 + 10;
class Solution {
    public:
    int constrainedSubsetSum(vector<int>& v, int k) {
        int ret = -inf;
        vector<int> dp(v.begin(), v.end());
        deque<int> dq;
        dq.push_front(v[0]);
        int n = v.size();
        ret = v[0];
        for (int i = 1; i < n; i++) {
            if (i > k && dq.front() == dp[i - k - 1])
            dq.pop_front();
            dp[i] = max(dp[i], dq.empty() ? dp[i] + 0 : dp[i] +
            dq.front());
            while (!dq.empty() && dq.back() < dp[i])
            dq.pop_back();
            dq.push_back(dp[i]);
            ret = max(ret, dp[i]);
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {10,2,-9,5,19};
    cout << (ob.constrainedSubsetSum(v, 2));
}

実行結果

入力:

{10,2,-9,5,19}, 2

出力:

36

計算量

各要素はデックに対して高々1回の追加と削除しか行われないため、時間計算量は O(n) です。また、dp 配列とデックの分だけメモリが必要となるため、空間計算量も O(n) となります。素朴な DP(O(nk))と比べて、大きな入力でも高速に動作する点がこの手法の大きな利点です。