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

C++에서 합이 K와 같은 부분 배열(Subarray)의 개수 구하기

정수 배열 nums와 정수 k가 주어졌을 때, 원소들의 합이 정확히 k와 같은 연속된 부분 배열(subarray)의 총 개수를 구하는 문제입니다. 예를 들어 배열이 [1, 1, 1]이고 k가 2라면, 답은 2가 됩니다. 첫 두 원소로 이루어진 [1, 1]과 마지막 두 원소로 이루어진 [1, 1], 이렇게 두 가지 경우가 해당되기 때문입니다.

접근 방법 — 누적 합(Prefix Sum)과 해시 맵

모든 부분 배열을 일일이 확인하는 브루트 포스 방식은 O(n²)의 시간이 소요됩니다. 누적 합과 해시 맵을 함께 사용하면 이를 O(n)으로 최적화할 수 있습니다.

핵심 아이디어는 다음과 같습니다. 인덱스 j까지의 누적 합을 P(j)라고 하면, 구간 [j+1, i]의 합은 P(i) − P(j)입니다. 따라서 P(i) − P(j) = k, 즉 P(j) = P(i) − k를 만족하는 과거 지점 j의 개수만 세면 됩니다. 참고로 이 구현에서는 맵에 누적 합의 음수 값(−temp)을 저장하기 때문에 조회 키가 k − temp가 된다는 점에 유의하세요.

알고리즘 단계

  1. 누적 합의 등장 횟수를 저장할 맵 sums를 준비하고, temp := 0, ans := 0으로 초기화한 뒤 sums[0] := 1을 설정합니다. (아무 원소도 선택하지 않은 빈 접두사를 의미)
  2. 배열의 각 원소 n[i]에 대해 다음을 반복합니다.
    • temp := temp + n[i] 로 현재까지의 누적 합을 갱신합니다.
    • 맵에 k − temp 키가 존재하면, ans := ans + sums[k − temp] 를 수행합니다. 해당 횟수만큼 합이 k인 부분 배열을 새로 찾은 것입니다.
    • sums[−temp] 값을 1 증가시켜 현재 누적 합을 기록합니다.
  3. 모든 순회가 끝나면 ans를 반환합니다.

C++ 구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int subarraySum(vector<int>& n, int k) {
      unordered_map <int, int> sums;
      int temp = 0;
      sums[0] = 1;
      int ans =0;
      for(int i =0;i<n.size();i++){
         temp+= n[i];
         if(sums.find(k-temp)!=sums.end()){
            ans += sums[k-temp];
         }
         sums[-temp]++;
      }
      return ans;
   }
};
main(){
   Solution ob;
   vector<int> v = {1,1,1};
   cout << (ob.subarraySum(v, 2));
}

입력 및 출력

입력:

[1,1,1]
2

출력:

2

예제 실행 과정 살펴보기

[1, 1, 1], k = 2인 경우를 단계별로 추적하면 다음과 같습니다.

  • 초기 상태: sums = {0: 1}, temp = 0, ans = 0
  • i = 0: temp = 1, k − temp = 1 → 맵에 없음. sums[−1] = 1 기록
  • i = 1: temp = 2, k − temp = 0 → 맵에 있음! ans += 1 → ans = 1. sums[−2] = 1 기록
  • i = 2: temp = 3, k − temp = −1 → 맵에 있음! ans += 1 → ans = 2. sums[−3] = 1 기록

최종적으로 ans = 2가 반환됩니다.

복잡도 분석

  • 시간 복잡도: O(n) — 배열을 한 번만 순회하며, 해시 맵의 삽입·조회는 평균적으로 O(1)입니다.
  • 공간 복잡도: O(n) — 최악의 경우 서로 다른 누적 합이 n개까지 저장될 수 있습니다.