C++로 풀어보는 범위 합계 개수(Range Sum Count) 문제

문제 소개

정수 배열 nums가 주어졌다고 가정해 봅시다. 우리가 구해야 할 것은 [lower, upper] 범위(양쪽 경계 포함)에 속하는 구간 합(range sum)의 개수입니다. 여기서 구간 합 S(i, j)는 인덱스 i부터 j까지(i ≤ j)에 해당하는 요소들의 합으로 정의됩니다.

예를 들어 입력 배열이 [-3, 6, -1]이고 lower = -2, upper = 2라고 해 보겠습니다. 이때 조건을 만족하는 구간은 [0, 2](합 = 2)와 [2, 2](합 = -2) 두 가지이므로 결과값은 2가 됩니다.

접근 방법

이 문제는 누적 합(prefix sum)과 병합 정렬(merge sort)을 결합하면 O(n log n) 시간 복잡도로 효율적으로 해결할 수 있습니다. 구간 합 S(i, j) = prefix[j+1] − prefix[i]이므로, 결국 “prefix 배열에서 두 원소의 차이가 [lower, upper] 범위 안에 들어오는 쌍의 개수”를 세는 문제로 바꿔 생각할 수 있습니다.

알고리즘 단계

mergeIt() 함수를 정의합니다. 이 함수는 prefix 배열과 start, mid, end, lower, upper 값을 인자로 받습니다.
i := start, j := mid + 1로 초기화합니다.
temp := end − start + 1 (임시 배열의 크기)
low := mid + 1, high := mid + 1로 초기화합니다.
k := 0으로 초기화합니다.
크기가 temp인 배열 arr을 선언합니다.
i ≤ mid인 동안 다음을 반복합니다.
- low ≤ end이면서 prefix[low] − prefix[i] < lower인 동안 low를 1씩 증가시킵니다.
- high ≤ end이면서 prefix[high] − prefix[i] ≤ upper인 동안 high를 1씩 증가시킵니다.
- j ≤ end이면서 prefix[j] < prefix[i]인 동안 arr[k] := prefix[j]를 저장하고 j와 k를 각각 1씩 증가시킵니다.
- arr[k] := prefix[i]를 저장한 뒤 i와 k를 1씩 증가시킵니다.
- count := count + (high − low)로 조건을 만족하는 쌍의 개수를 누적합니다.
j ≤ end인 동안 오른쪽 절반에 남은 원소들을 arr에 차례대로 복사합니다.
arr의 내용을 prefix[start..end] 구간에 다시 복사하여 정렬된 상태를 유지합니다.

merge() 함수

start ≥ end이면 더 이상 나눌 수 없으므로 함수를 종료합니다.
mid := start + (end − start) / 2로 분할 지점을 계산합니다.
왼쪽 절반에 대해 merge(prefix, start, mid, lower, upper)를 재귀 호출합니다.
오른쪽 절반에 대해 merge(prefix, mid + 1, end, lower, upper)를 재귀 호출합니다.
mergeIt(prefix, start, mid, end, lower, upper)를 호출해 두 구간을 병합하면서 조건을 만족하는 쌍을 셉니다.

메인 함수 처리 과정

n := nums의 크기, count := 0으로 초기화합니다.
크기가 n + 1인 prefix 배열을 선언하고 prefix[0] := 0으로 설정합니다.
i = 1부터 n까지 prefix[i] := prefix[i−1] + nums[i−1]로 누적 합을 계산합니다.
merge(prefix, 0, n, lower, upper)를 호출합니다.
count 값을 반환합니다.

아래 구현 예시를 통해 더 자세히 이해해 보겠습니다.

예제 코드 (C++)

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
   int count = 0;
   void mergeIt(lli prefix[], lli start ,lli mid, lli end, lli lower, lli upper){
      lli i = start, j = mid + 1;
      lli temp = end - start + 1;
      lli low = mid + 1, high = mid + 1;
      lli k = 0;
      lli arr[temp];
      while(i <= mid){
         while(low <= end && prefix[low] - prefix[i] < lower) low++;
         while(high <= end && prefix[high] - prefix[i] <= upper) high++;
         while(j<= end && prefix[j] < prefix[i]){
            arr[k] = prefix[j];
            j++;
            k++;
         }
         arr[k] = prefix[i];
         i++;
         k++;
         count += high - low;
      }
      while(j <= end){
         arr[k] = prefix[j];
         k++;
         j++;
      }
      for(i = 0; i < temp; i++){
         prefix[start] = arr[i];
         start++;
      }
   }
   void merge(lli prefix[], lli start, lli end, lli lower, lli upper){
      if(start >= end)return;
      lli mid = start + (end - start) / 2;
      merge(prefix, start, mid, lower, upper);
      merge(prefix, mid + 1, end, lower, upper);
      mergeIt(prefix, start, mid, end, lower, upper);
   }
   int countRangeSum(vector<int>& nums, int lower, int upper) {
      lli n = nums.size();
      count = 0;
      lli prefix[n + 1];
      prefix[0] = 0;
      for(lli i = 1; i <= n; i++){
         prefix[i] = prefix[i - 1] + nums[i - 1];
      }
      merge(prefix, 0, n, lower, upper);
      return count;
   }
};
main(){
   Solution ob;
   vector<int> v = {-3,6,-1};
   cout << (ob.countRangeSum(v, -2, 2));
}

입력

{-3,6,-1}
-2
2

출력

복잡도 분석

시간 복잡도: O(n log n) — 병합 정렬의 분할과 병합 단계마다 선형 시간의 탐색이 수행되기 때문입니다. 완전 탐색으로 모든 구간을 확인하는 O(n²) 방식보다 훨씬 효율적입니다.

공간 복잡도: O(n) — 누적 합 배열과 병합 과정에서 사용되는 임시 배열을 위해 추가 공간이 필요합니다.