C++で範囲[L、R]内における最大K回の移動での数値の合計を最大化する

本記事では、整数を含む配列 Arr[] と、複数のクエリを格納した2次元配列 Q が与えられたときの問題を解説します。各クエリは lpos（開始インデックス）、rpos（終了インデックス）、K（最大ステップ数） の3つの値で構成されています。

インデックスiからは、1ステップで次のインデックスi+1へ進むか、現在の位置にとどまることができます。lposからrposへは、最大Kステップ以内で移動しなければなりません。移動中は、左端の数値を含めて各ステップで訪れた位置の数値をすべて加算し、その合計を最大化することが目標です。Kステップ以内でlposからrposへの移動が不可能な場合は「No」を出力します。それでは詳しく見ていきましょう。

入出力シナリオ

入力 − Arr[] = {1, 2, 4, -1 };

Q[][3] = { { 0, 2, 2 }, { 0, 2, 1 }, { 3, 3, 1 }, { 0, 2, 3} };

出力 −

クエリ1: 7
クエリ2: NO
クエリ3: NO
クエリ4: 11

説明 −

最初のクエリ

最大2ステップでインデックス0から2へ移動できます。

ステップ1：インデックス0から1へ（1+2=3）
ステップ2：インデックス1から2へ（3+4=7）

2番目のクエリ

最大1ステップではインデックス0から2へ移動できません。「NO」を出力します。

3番目のクエリ

最大1ステップではインデックス3から3へ移動できません。「NO」を出力します。

4番目のクエリ

最大3ステップでインデックス0から2へ移動できます。

ステップ1：インデックス0から1へ（1+2=3）
ステップ2：インデックス1から2へ（3+4=7）
ステップ3：インデックス2にとどまる（7+4=11）

入力 − Arr[] = { 1, 2, 3, 3, 2 }; Q[][3] = { { 0, 3, 2 }, { 1, 4, 3 } };

出力 −

クエリ1: NO
クエリ2: 10

説明 −

最初のクエリ

最大2ステップではインデックス0から3へ移動できません。「NO」を出力します。

2番目のクエリ

最大3ステップでインデックス1から4へ移動できます。

ステップ1：インデックス1から2へ（2+3=5）
ステップ2：インデックス2から3へ（5+3=8）
ステップ3：インデックス3から4へ（8+2=10）

プログラムで使用しているアプローチ

このアプローチでは、セグメント木を使用して範囲lpos〜rpos内の最大値を求め、累積和（プレフィックスサム）を使用してすべての数値の合計を計算します。

• 入力配列Arr[]とクエリ行列Q[][]を受け取ります。

• セグメント木を実装するための配列 sgTreee[5 * length] を用意します。

• 累積和用の配列 pSum[length] を用意します。

• 関数 createTree(int min, int max, int pos, int sgT[], int arr[], int len) は、セグメント木のノード値を構築します。

• (min == max) の場合は葉ノードであることを意味するため、sgT[pos] = arr[max] を設定します。

• midd = (min + max) / 2 を求めます。

• 左右の子部分木に対して createTree(min, midd, loc1, sgT, arr, len) と createTree(midd + 1, max, loc2, sgT, arr, len) を再帰的に呼び出します。ここで loc1=2*pos+1、loc2=2*pos+2 です。

• tmp1=sgT[loc1]、tmp2=sgT[loc2] を取得し、sgT[pos] に両者のうち大きい方を設定します。

• 関数 preSum(int pSum4[], int arr4[], int len4) は入力配列を受け取り、forループを使って累積和配列を更新します。

• インデックス1から末尾までの各要素に対して、pSum4[j] = pSum4[j - 1] + arr4[j]; と更新します。

• 関数 resQuery(int len3, int arr3[], int sgT3[], int pSum3[], int q1[][3], int qlen1) はすべての入力パラメータを受け取り、各クエリの結果を出力します。

• resQuery() の内部では、forループを使って solQuery(int lpos, int rpos, int k, int len2, int arr2[], int sgT2[], int pSum2[]) を呼び出し、各クエリを順番に処理します。

• 関数 solQuery() はクエリを解いて結果を返します。

• rpos - lpos > k の場合は解が存在しないため、-1 を返します。

• maxVal = findMax(0, len2 - 1, lpos, rpos, 0, sgT2, arr2, len2); として範囲内の最大値を取得します。

• maxVal < 0 の場合は maxVal を0に設定します。

• 変数 sum = pSum2[rpos] とします。

• lpos > 0 の場合は sum -= pSum2[lpos - 1] とし、result = sum + (k - (rpos - lpos)) * maxVal を計算します。

• result を返します。

• 関数 findMax(int start, int end, int min1, int max1, int pos1, int sgT1[], int arr1[], int len1) は、範囲lpos〜rpos間の最大値を返します。

• (min1 <= start) かつ (max1 >= end) の場合は完全に重なっているため、sgT1[pos1] を返します。

• (end < min1 || start > max1) の場合は範囲外なので INT_MIN を返します。

• 左右の子部分木に対する再帰呼び出しで lmax と rmax を計算し、両者の最大値を返します。

• 最後に各クエリの結果が出力されます。解が存在しない場合は「No」と表示されます。

例

#include <bits/stdc++.h>
using namespace std;
void createTree(int min, int max, int pos,
int sgT[], int arr[], int len){ if (min == max) {
    sgT[pos] = arr[max];
    return;
    }
    int midd = (min + max) / 2;
    int loc1=2*pos+1;
    int loc2=2*pos+2;
    createTree(min, midd, loc1, sgT, arr, len);
    createTree(midd + 1, max, loc2, sgT, arr, len);
    int tmp1=sgT[loc1];
    int tmp2=sgT[loc2];
    sgT[pos] = tmp1>tmp2 ? tmp1 : tmp2 ;
}
int findMax(int start, int end, int min1, int max1, int pos1, int sgT1[], int arr1[], int len1){
    int middle;
    if (min1 <= start)
    { if( max1 >= end){
            return sgT1[pos1];
      }
    }
    if (end < min1 || start > max1)
    { return INT_MIN; }

    middle = (start + end) / 2;
    int loc1=2 * pos1 + 1;
    int loc2=2 * pos1 + 2;
    int lmax = findMax(start, middle, min1, max1, loc1, sgT1, arr1, len1);
    int rmax = findMax(middle + 1, end, min1, max1, loc2, sgT1, arr1, len1);
    int res=lmax>rmax?lmax:rmax;
    return res;
}
int solQuery(int lpos, int rpos, int k, int len2, int arr2[], int sgT2[], int pSum2[]){
    int result;
       if (rpos - lpos > k)
       { return -1; }
       int maxVal = findMax(0, len2 - 1, lpos, rpos, 0, sgT2, arr2, len2);
       if (maxVal < 0)
       { maxVal = 0; }
       int sum = pSum2[rpos];
       if (lpos > 0)
       { sum -= pSum2[lpos - 1]; }
       result = sum + (k - (rpos - lpos)) * maxVal;
       return result;
    }
    void resQuery(int len3, int arr3[], int sgT3[],
          int pSum3[], int q1[][3], int qlen1){
       int i;
       int result;
       for (i = 0; i < qlen1; i++) {
       result = solQuery(q1[i][0], q1[i][1],q1[i][2], len3, arr3, sgT3, pSum3);

       if (result == -1)
          { cout <<endl<<"Query "<<i+1<<": "<<"NO"; }
       else
          { cout <<endl<<"Query "<<i+1<<": "<<result; }
       }
    }
void preSum(int pSum4[], int arr4[], int len4){
    pSum4[0] = arr4[0];
    int j;
    for (j = 1; j < len4; j++){
       pSum4[j] = pSum4[j - 1] + arr4[j];
    }
}
int main(){
   int Arr[] = {1, 2, 4, -1 };
    int length = sizeof(Arr) / sizeof(Arr[0]);
    int sgTreee[5 * length];
    createTree(0, length - 1, 0, sgTreee, Arr, length);
   int pSum[length];
   preSum(pSum, Arr, length);
   int Q[][3] = { { 0, 2, 2 },
      { 0, 2, 1 },
      { 3, 3, 1 },
      { 0, 2, 3} };
   int qlen = sizeof(Q) / sizeof(Q[0]);
   resQuery(length, Arr, sgTreee, pSum, Q, qlen);
   return 0;
}

出力

上記のコードを実行すると、以下の出力が生成されます。

Query 1: 7
Query 2: NO
Query 3: NO
Query 4: 11