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

C++에서 합이 같은 네 부분으로 배열 분할하는 방법

문제 개요

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)의 합이 서로 동일

여기서 부분 배열 (L, R)은 원본 배열에서 인덱스 L부터 R까지의 구간을 의미합니다. 즉, 배열을 네 개의 연속된 구간으로 나누었을 때 각 구간의 합이 모두 같은지를 확인하는 문제입니다.

예시

입력이 [1,2,1,2,1,2,1]이라면 결과는 True입니다. i = 1, j = 3, k = 5로 선택하면 네 구간의 합이 각각 1로 동일하기 때문입니다.

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

풀이 접근 방식

모든 (i, j, k) 조합을 무작정 탐색하면 시간이 너무 오래 걸립니다. 대신 누적합(prefix sum)과 집합(set)을 활용하면 효율적으로 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다.

  1. 먼저 누적합 배열 sums를 계산합니다. sums[i]는 nums[0]부터 nums[i]까지의 합을 담습니다.
  2. 중간 분할점 j를 하나씩 이동시키며, 왼쪽 두 구간의 합이 같아지는 i 값을 찾아 그 합을 집합 s에 저장합니다.
  3. 오른쪽 두 구간에 대해서는 k를 이동시키며 두 구간의 합이 같은지 확인하고, 그 값이 집합 s에 존재하는지 검사합니다. 네 구간의 합이 모두 같으려면 왼쪽 두 구간의 공통 합과 오른쪽 두 구간의 공통 합이 일치해야 하기 때문입니다.

알고리즘 단계

  • n := nums의 크기로 설정
  • 크기가 n인 배열 sums 정의
  • sums[0] := nums[0]
  • i := 1부터 n-1까지 반복하며 sums[i] := nums[i] + sums[i-1]로 누적합 계산
  • j := 3부터 n-1까지 반복:
    • 집합 s 생성
    • i := 1부터 j-2까지 반복:
      • sum1 := sums[i-1]
      • sum2 := sums[j-1] - sums[i]
      • sum1 == sum2이면 sum1을 s에 삽입
    • k := j+2부터 n-2까지 반복:
      • sum1 := sums[k-1] - sums[j]
      • sum2 := sums[n-1] - sums[k]
      • sum1 == sum2이고 sum1이 s에 존재하면 true 반환
  • 모든 경우를 확인한 후에도 찾지 못하면 false 반환

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²)입니다. 중간 분할점 j마다 왼쪽 탐색과 오른쪽 탐색을 각각 수행하므로 전체적으로 이중 반복문 형태가 되며, 단순 브루트 포스 방식의 O(n³)보다 훨씬 효율적입니다. 공간 복잡도는 누적합 배열과 집합 저장에 사용되는 O(n)입니다.