Javaで指定された分割点に基づき配列を部分配列に分割した際の最大部分配列和を求める方法

問題の概要

2つの整数配列が与えられます。1つは合計を計算する対象となる要素を含む配列、もう1つは配列を部分配列（部分集合）へ分割するための分割点を含む配列です。分割が行われるたびに、その時点で存在するすべての部分配列の合計を計算し、最大の合計値を出力します。

例で理解する

入力 − int arr[] = { 9, 4, 5, 6, 7 }、int splitPoints[] = { 0, 2, 3, 1 }

出力 − 各分割後の最大部分配列和：[22, 13, 9, 9]

説明 − 配列を分割点に従って分割し、各段階での最大部分配列和を求めます。

• 1回目の分割後 → {9} と {4,5,6,7} ／ 最大部分配列和は 22
• 2回目の分割後 → {9}、{4,5}、{6,7} ／ 最大部分配列和は 13
• 3回目の分割後 → {9}、{4,5}、{6}、{7} ／ 最大部分配列和は 9
• 4回目の分割後 → {9}、{4}、{5}、{6}、{7} ／ 最大部分配列和は 9

入力 − int arr[] = { 7, 8, 5, 9, 1 }、int splitPoints[] = { 1, 2, 0, 3 }

出力 − 各分割後の最大部分配列和：[15, 15, 10, 9]

説明 − 同様に、配列を分割点に従って分割し、各段階での最大部分配列和を求めます。

• 1回目の分割後 → {7,8} と {5,9,1} ／ 最大部分配列和は 15
• 2回目の分割後 → {7,8}、{5}、{9,1} ／ 最大部分配列和は 15
• 3回目の分割後 → {7}、{8}、{5}、{9,1} ／ 最大部分配列和は 10
• 4回目の分割後 → {7}、{8}、{5}、{9}、{1} ／ 最大部分配列和は 9

プログラムで使用するアプローチ

この問題は、累積和（プレフィックスサム）と TreeSet を組み合わせることで効率的に解くことができます。手順は以下の通りです。

• main() メソッドから開始
  • 任意の長さの配列 arr[] と splitPoints[] を用意し、それぞれの長さを計算したうえで、calculateSubsetSum(arr.length, splitPoints.length, splitPoints, arr) を呼び出します。
• calculateSubsetSum() メソッド内
  • 整数配列 sum[] を作成し、sum[0] を arr[0] に設定します。
  • i = 1 から配列の長さまでループし、sum[i] = sum[i - 1] + arr[i] として累積和を構築します。また、temp[0] を新しい subSets(0, n - 1, sum[n - 1]) に設定します。
  • t2.add(temp[0]) および t1.add(0) を実行します。
  • i = 0 から splitPoints の長さまでループします。ループ内では currentSplitPoint を t1.floor(splitPoints[i]) として取得し、t2.remove(temp[currentSplitPoint]) で該当する部分配列を削除します。
  • end を temp[currentSplitPoint].last とし、temp[currentSplitPoint] を新しい subSets(currentSplitPoint, splitPoints[i], sum[splitPoints[i]] - (currentSplitPoint == 0 ? 0 : sum[currentSplitPoint - 1])) に更新します。
  • t2.add(temp[currentSplitPoint]) を実行し、さらに temp[splitPoints[i] + 1] を新しい subSets(splitPoints[i] + 1, end, sum[end] - sum[splitPoints[i]]) として作成します。
  • t2.add(temp[splitPoints[i] + 1])、t1.add(currentSplitPoint)、t1.add(splitPoints[i] + 1) を実行します。
  • t2.first() の値を出力します。
• subSets クラスの作成
  • first、last、value をデータメンバーとして宣言し、コンストラクタ subSets(int f, int l, int v) を定義して、first に f、last に l、value に v を設定します。
• utilityComparator クラスの作成（Comparator<subSets> を実装）
  • public メソッド compare を作成し、s2.value と s1.value が等しくない場合は s2.value - s1.value を返します。
  • s1.first と s2.first が等しくない場合は s2.first - s1.first を返します。

コード例

import java.io.IOException;
import java.io.InputStream;
import java.util.*;
class utilityComparator implements Comparator<subSets>{
    public int compare(subSets s1, subSets s2){
        if(s2.value != s1.value){
            return s2.value - s1.value;
        }
        if(s1.first != s2.first){
            return s2.first - s1.first;
        }
        return 0;
    }
}
class subSets{
    int first;
    int last;
    int value;
    subSets(int f, int l, int v){
        first = f;
        last = l;
        value = v;
    }
}
public class testClass{
    static void calculateSubsetSum(int n, int k, int splitPoints[], int arr[]){
        int sum[] = new int[n];
        sum[0] = arr[0];
        for (int i = 1; i < n; i++){
            sum[i] = sum[i - 1] + arr[i];
        }
        TreeSet<Integer> t1 = new TreeSet<>();
        TreeSet<subSets> t2 = new TreeSet<>(new utilityComparator());
        subSets temp[] = new subSets[n];
        temp[0] = new subSets(0, n - 1, sum[n - 1]);
        t2.add(temp[0]);
        t1.add(0);
        System.out.println("Maximum subarray sum after each split");
        for (int i = 0; i < k; i++){
            int currentSplitPoint = t1.floor(splitPoints[i]);
            t2.remove(temp[currentSplitPoint]);
            int end = temp[currentSplitPoint].last;
            temp[currentSplitPoint] = new subSets(currentSplitPoint, splitPoints[i], sum[splitPoints[i]] - (currentSplitPoint == 0 ? 0 : sum[currentSplitPoint - 1]));
            t2.add(temp[currentSplitPoint]);
            temp[splitPoints[i] + 1] = new subSets(splitPoints[i] + 1, end, sum[end] - sum[splitPoints[i]]);
            t2.add(temp[splitPoints[i] + 1]);
            t1.add(currentSplitPoint);
            t1.add(splitPoints[i] + 1);
            System.out.println(t2.first().value);
        }
    }
    public static void main(String[] args){
        int arr[] = { 2, 1, 6, 8, 5, 10, 21, 13};
        int splitPoints[] = { 3, 1, 2, 0, 4, 5 };
        calculateSubsetSum(arr.length, splitPoints.length, splitPoints, arr);
    }
}

出力結果

上記のコードを実行すると、次の出力が得られます。

Maximum subarray sum after each split
49
49
49
49
44
34

このように、累積和によって各部分配列の合計を O(1) で算出し、TreeSet によって最大値の管理や挿入・削除を O(log n) で行うことで、分割が繰り返されても全体の処理を効率的に保つことができます。