C++で解く「3つの重ならない部分配列の最大合計」問題：累積和を使った効率的な実装

正の整数からなる配列 nums が与えられたとき、合計が最大となる3つの重ならない部分配列を見つける問題を考えます。各部分配列の長さは k に固定されており、3 * k 個の要素全体の合計を最大化することが目標です。

答えは、各区間の開始位置を表すインデックスのリストとして返します。条件を満たす組み合わせが複数存在する場合は、辞書順で最小のものを返す必要があります。

たとえば、入力が [1,2,1,2,6,8,4,1]、k = 2 の場合、出力は [0,3,5] となります。これは、部分配列 [1,2]、[2,6]、[8,4] の開始インデックス [0,3,5] に対応しています。

アルゴリズムのアプローチ

この問題は、累積和（プレフィックスサム）と、左側・右側それぞれから見た「その時点までの最大合計を持つ区間の開始位置」を事前計算しておくことで、効率的に解くことができます。手順は以下の通りです。

• n := 配列 nums のサイズとする
• サイズ3の配列 ret を定義する（初期値は無限大）
• サイズ n + 1 の累積和配列 sum を定義する
• i := 0 から i < n まで繰り返す：
  • sum[i + 1] = sum[i] + nums[i]
• サイズ n の配列 posLeft を定義する
• サイズ n の配列 posRight を定義し、n - k で初期化する
• i := k、currMax := sum[k] - sum[0] として、i < n の間 i を1ずつ増やしながら処理する：
  • newTotal := sum[i + 1] - sum[i + 1 - k]
  • newTotal > currMax の場合：
    • currMax := newTotal
    • posLeft[i] := i + 1 - k
  • それ以外の場合：
    • posLeft[i] := posLeft[i - 1]
• i := n - k - 1、currMax := sum[n] - sum[n - k] として、i >= 0 の間 i を1ずつ減らしながら処理する：
  • newTotal := sum[i + k] - sum[i]
  • newTotal >= currMax の場合：
    • currMax := newTotal
    • posRight[i] := i
  • それ以外の場合：
    • posRight[i] := posRight[i + 1]
• req := 0 とする
• i := k から i <= n - 2 * k まで繰り返す：
  • l := posLeft[i - 1]、r := posRight[i + k]
  • temp := (sum[l + k] - sum[l]) + (sum[i + k] - sum[i]) + (sum[r + k] - sum[r])
  • temp > req の場合：
    • ret[0] := l、ret[1] := i、ret[2] := r
    • req := temp
• ret を返す

ポイント解説

posLeft[i] は「インデックス i より前（左側）で、長さ k の部分配列として合計が最大になる開始位置」を保持します。同様に posRight[i] は「インデックス i 以降（右側）で最大合計となる開始位置」を保持します。中央の区間の開始位置 i を全探索しながら、左右の最適な区間と組み合わせることで、全体の最大合計を求められます。

なお、右側の走査では比較に >= を使うことで、同点の場合により小さいインデックスが選ばれ、結果として辞書順最小の解が得られます。

C++による実装例

以下のコードで実際の動作を確認してみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
class Solution {
public:
   vector<int> maxSumOfThreeSubarrays(vector<int>& nums, int k) {
      int n = nums.size();
      vector <int> ret(3, INT_MAX);
      vector <int> sum(n + 1);
      for(int i = 0; i < n; i++){
         sum[i + 1] = sum[i] + nums[i];
      }
      vector <int> posLeft(n);
      vector <int> posRight(n, n - k);
      for(int i = k, currMax = sum[k] - sum[0]; i < n; i++){
         int newTotal = sum[i + 1] - sum[i + 1- k];
         if(newTotal > currMax){
            currMax = newTotal;
            posLeft[i] = i + 1 - k;
         }else{
            posLeft[i] = posLeft[i - 1];
         }
      }
      for(int i = n - k - 1, currMax = sum[n] - sum[n - k]; i >=0 ; i--){
         int newTotal = sum[i + k] - sum[i];
         if(newTotal >= currMax){
            currMax = newTotal;
            posRight[i] = i;
         }else{
            posRight[i] = posRight[i + 1];
         }
      }
      int req = 0;
      for(int i = k; i <= n - 2 * k; i++){
         int l = posLeft[i - 1];
         int r = posRight[i + k];
         int temp = (sum[l + k] - sum[l]) + (sum[i + k] - sum[i]) + (sum[r + k] - sum[r]);
         if(temp > req){
            ret[0] = l;
            ret[1] = i;
            ret[2] = r;
            req = temp;
         }
      }
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v = {1,2,1,2,6,8,4,1};
   print_vector(ob.maxSumOfThreeSubarrays(v, 2));
}

入力

{1,2,1,2,6,8,4,1}
2

出力

[0, 3, 5]

まとめ

この手法では、累積和の構築に O(n)、左右の前計算にそれぞれ O(n)、最後の結合処理にも O(n) しかかからないため、全体の計算量は O(n) となり、非常に効率的です。素朴な全探索では O(n³) かかるところを、前計算によって大幅に高速化できる好例といえます。