Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 [L, R] 범위에서 최대 K번 이동 시 얻을 수 있는 숫자 합 최대화하기

정수를 담고 있는 배열 Arr[]과 쿼리를 담고 있는 2차원 배열 Q가 주어집니다. 각 쿼리는 세 개의 값, 즉 lpos(시작 위치), rpos(끝 위치), K(최대 이동 횟수)를 포함합니다.

한 번의 스텝으로 인덱스 i에서 다음 인덱스 i+1로 이동하거나, 현재 인덱스에 머무를 수 있습니다. lpos에서 rpos까지는 최대 K번의 스텝 안에서만 이동할 수 있으며, 가장 왼쪽 숫자를 포함해 각 스텝마다 방문한 숫자를 모두 더합니다. 목표는 최대 K번의 이동 안에서 합을 최대화하는 것입니다. 만약 K번의 스텝으로 lpos에서 rpos까지 이동하는 것이 불가능하다면 "NO"를 출력해야 합니다.

입출력 예시 살펴보기

입력 − Arr[] = {1, 2, 4, -1};
Q[][3] = { { 0, 2, 2 }, { 0, 2, 1 }, { 3, 3, 1 }, { 0, 2, 3 } };

출력 −

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

첫 번째 쿼리 설명

인덱스 0에서 2까지 최대 2번의 스텝으로 이동할 수 있습니다.

  • Step 1: 인덱스 0 → 1 (1 + 2 = 3)
  • Step 2: 인덱스 1 → 2 (3 + 4 = 7)

두 번째 쿼리 설명

인덱스 0에서 2까지는 거리가 2이므로 최대 1번의 스텝으로는 도달할 수 없습니다. 따라서 "NO"를 출력합니다.

세 번째 쿼리 설명

인덱스 3에서 3으로 이동하려면 시작점과 끝점이 같아 이동 자체가 불가능합니다. 따라서 "NO"를 출력합니다.

네 번째 쿼리 설명

인덱스 0에서 2까지 최대 3번의 스텝으로 이동할 수 있습니다.

  • Step 1: 인덱스 0 → 1 (1 + 2 = 3)
  • Step 2: 인덱스 1 → 2 (3 + 4 = 7)
  • Step 3: 인덱스 2에 머무름 (7 + 4 = 11)

남은 스텝 동안 최댓값이 있는 위치에 머물면서 합을 추가로 늘릴 수 있다는 점이 핵심입니다.

두 번째 입출력 예시

입력 − Arr[] = { 1, 2, 3, 3, 2 }; Q[][3] = { { 0, 3, 2 }, { 1, 4, 3 } };

출력 −

Query 1: NO
Query 2: 10

첫 번째 쿼리 설명

인덱스 0에서 3까지는 거리가 3이므로 최대 2번의 스텝으로는 도달할 수 없습니다. "NO"를 출력합니다.

두 번째 쿼리 설명

인덱스 1에서 4까지 최대 3번의 스텝으로 이동할 수 있습니다.

  • Step 1: 인덱스 1 → 2 (2 + 3 = 5)
  • Step 2: 인덱스 2 → 3 (5 + 3 = 8)
  • Step 3: 인덱스 3 → 4 (8 + 2 = 10)

풀이 접근 방식

이 문제는 세그먼트 트리(Segment Tree)를 활용해 lpos부터 rpos 범위 내의 최댓값을 빠르게 찾고, 접두사 합(Prefix Sum) 배열을 사용해 구간 합을 계산하는 방식으로 해결할 수 있습니다.

  • 입력 배열 Arr[]과 쿼리 행렬 Q[][]를 받습니다.
  • 세그먼트 트리 구현을 위해 크기가 5 * length인 sgTree[] 배열을 선언합니다.
  • 접두사 합을 저장할 pSum[length] 배열을 준비합니다.

createTree 함수 — 세그먼트 트리 생성

  • 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, ...)와 createTree(midd + 1, max, loc2, ...)를 재귀 호출합니다. 여기서 loc1 = 2*pos+1, loc2 = 2*pos+2입니다.
  • tmp1 = sgT[loc1], tmp2 = sgT[loc2]를 구한 뒤, 둘 중 더 큰 값으로 sgT[pos]를 갱신합니다.

preSum 함수 — 접두사 합 계산

  • preSum(int pSum4[], int arr4[], int len4) 함수는 입력 배열을 받아 for 루프로 접두사 합 배열을 갱신합니다.
  • 인덱스 1부터 마지막까지 pSum4[j] = pSum4[j - 1] + arr4[j]로 업데이트합니다.

resQuery 함수 — 쿼리 결과 처리

  • resQuery(int len3, int arr3[], int sgT3[], int pSum3[], int q1[][3], int qlen1) 함수는 모든 입력값을 받아 각 쿼리의 결과를 출력합니다.
  • 내부적으로 solQuery()를 호출하여 for 루프로 각 쿼리를 하나씩 해결합니다.

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입니다. 즉, 남는 스텝만큼 최댓값을 반복해서 더합니다.

findMax 함수 — 범위 내 최댓값 탐색

  • 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