수들의합4.md

DONE_DATE	2025/01/22

수들의 합 4

난이도

• 골드 4

문제

https://www.acmicpc.net/problem/2015

알고리즘 분류

• 자료 구조
• 누적 합
• 해시를 사용한 집합과 맵
• 트리를 사용한 집합과 맵

회고

N과 A[1], A[2], ..., A[N]이 주어졌을 때, 이러한 N×(N+1)/2개의 부분합 중 합이 K인 것이 몇 개나 있는지 구하는 문제이다. 누적합을 쓰면 되지만 그럼 시간복잡도가 O(N^2)이 되어 시간초과가 날 것이다. 그래서 Map 을 이용해서 풀 수 있다.

#include <bits/stdc++.h>

typedef long long ll;
using namespace std;

int main() {
    int N, K;
    cin >> N >> K;
    vector<int> prefixSum(N + 1, 0);
    unordered_map<ll, ll> map;
    for (int i = 1; i <= N; i++) {
        int num;
        cin >> num;
        prefixSum[i] = prefixSum[i - 1] + num;
    }

누적합을 구하고, map을 선언해준다.

    
    
    ll ans = 0;
    for (int i = 1; i <= N; i++) {
        if (prefixSum[i] == K) {
            ans++;
        }
        ans += map[prefixSum[i] - K];
        map[prefixSum[i]]++;
    }
    cout << ans;
}

현재 누적 합 prefixSum[i] 자체가 K와 같다면, 1개의 부분 배열이 존재함을 의미하므로 ans를 증가시킵니다. prefixSum[i] - K가 이전에 몇 번 등장했는지를 map에서 조회합니다. 이는 현재 위치 i에서 j 까지의 부분 배열 합이 K인 경우를 의미합니다. 이 값을 ans에 더해줍니다. 현재 누적 합 prefixSum[i]을 해시 맵에 기록하여 이후에 이를 기반으로 한 계산이 가능하도록 합니다.

전체코드

#include <bits/stdc++.h>

typedef long long ll;
using namespace std;

int main() {
    int N, K;
    cin >> N >> K;
    vector<int> prefixSum(N + 1, 0);
    unordered_map<ll, ll> map;
    for (int i = 1; i <= N; i++) {
        int num;
        cin >> num;
        prefixSum[i] = prefixSum[i - 1] + num;
    }
    ll ans = 0;
    for (int i = 1; i <= N; i++) {
        if (prefixSum[i] == K) {
            ans++;
        }
        ans += map[prefixSum[i] - K];
        map[prefixSum[i]]++;
    }
    cout << ans;
}