【C++】配列を合計が等しい4つの部分配列に分割できるか判定する方法

問題概要

n個の整数からなる配列が与えられたとき、以下の条件を満たす三つ組 (i, j, k) が存在するかどうかを判定する問題を考えてみましょう。

• 0 < i、i + 1 < j、j + 1 < k < n − 1 を満たすこと

• 部分配列 (0, i−1)、(i+1, j−1)、(j+1, k−1)、(k+1, n−1) の4つの合計がすべて等しいこと

ここで、部分配列 (L, R) とは、元の配列のインデックス L の要素からインデックス R の要素までを切り出した範囲を意味します。

例として、入力が [1,2,1,2,1,2,1] の場合を考えてみます。このとき i = 1、j = 3、k = 5 となるため、出力は True になります。

sum(0, i − 1) = 1 → sum(0, 0) = 1
sum(i + 1, j − 1) = 1 → sum(2, 2) = 1
sum(j + 1, k − 1) = 1 → sum(4, 4) = 1
sum(k + 1, n − 1) = 1 → sum(6, 6) = 1

解法のアプローチ

この問題は、累積和（プレフィックスサム）とハッシュセットを組み合わせることで効率的に解くことができます。全要素の合計を前計算しておけば、任意の区間の合計を O(1) で求められるのがポイントです。

アルゴリズムの手順

• n := 配列 nums のサイズとする

• サイズ n の累積和配列 sums を定義する

• sums[0] := nums[0] と初期化する

• i = 1 から n 未満までループし、sums[i] := nums[i] + sums[i−1] として累積和を構築する

• j = 3 から n 未満までループする：

  • 整数を格納するセット s を用意する

  • i = 1 から j − 2 までループする：

    • sum1 := sums[i−1]（最初の区間の合計）

    • sum2 := sums[j−1] − sums[i]（2番目の区間の合計）

    • sum1 と sum2 が等しければ、sum1 をセット s に挿入する

  • k = j + 2 から n − 2 までループする：

    • sum1 := sums[k−1] − sums[j]（3番目の区間の合計）

    • sum2 := sums[n−1] − sums[k]（4番目の区間の合計）

    • sum1 と sum2 が等しく、かつ sum1 がセット s に存在すれば true を返す

• 条件を満たす組み合わせが見つからなければ false を返す

中央の分割点 j を固定した際、左側で成立する「最初の2区間の合計」をセットに記録しておき、右側の2区間の合計がその値と一致するかを確認することで、全体の計算量を O(n²) に抑えられます。

C++による実装例

それでは、実際のコードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   bool splitArray(vector<int>& nums) {
      int n = nums.size();
      vector<int> sums(n);
      sums[0] = nums[0];
      for (int i = 1; i < n; i++) {
         sums[i] += (nums[i] + sums[i - 1]);
      }
      for (int j = 3; j < n; j++) {
         set<int> s;
         for (int i = 1; i < j - 1; i++) {
            int sum1 = sums[i - 1];
            int sum2 = sums[j - 1] - sums[i];
            if (sum1 == sum2)
               s.insert(sum1);
         }
         for (int k = j + 2; k < n - 1; k++) {
            int sum1 = sums[k - 1] - sums[j];
            int sum2 = sums[n - 1] - sums[k];
            if (sum1 == sum2 && s.count(sum1))
               return true;
          }
         }
         return false;
      }
};
main(){
   Solution ob;
   vector<int> v = {1,2,1,2,1,2,1};
   cout << (ob.splitArray(v));
}

入力例

{1,2,1,2,1,2,1}

出力例

1

まとめ

このように、累積和を事前に計算しておくことで各区間の合計を高速に取得でき、セットを併用することで左右の候補を効率よく照合できます。全探索では O(n³) 以上かかるところを、O(n²) まで削減できる点がこの手法の大きな利点です。配列の分割に関する問題は競技プログラミングやコーディング面接でも頻出のテーマなので、ぜひパターンとして押さえておきましょう。