백준 2015번: 수들의 합 4 (Java)

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

2015번: 수들의 합 4

 

---

 

솔직히 누적합 유형 문제는 껌이라고 생각했는데, O(N^2) 브루트포스 풀이말고 생각이 전혀 안 나더라... 내가 껌이였고

 

풀이

1. 누적합(sum[i]: 0번째 index부터 i번째 index까지의 합)을 구한다.

2. value는 Integer, key는 long으로 하는 map을 선언한 후, (0, 1)을 넣어준다. 이는 0부터 i까지의 누적합이 K인 경우를 세주기 위함이다.

2. i = 0부터 N까지 반복문을 돌면서

 1) map에 sum[i] - K가 있는지 확인한다

  (1) 만약 sum[i]-K == 0이라면, 이는 0부터 i까지의 부분합이 우리가 원하는 부분합이 된다.

  (2) 다른 경우, 2-3) 과정에 의해 map에는 sum[0], ... , sum[i-1]이 있다. 따라서, 적당한 m < i가 존재하여 sum[i]-K = sum[m]이며, 이는 sum[i] - sum[m] = K, 즉 m+1번째 index부터 i번째 index까지의 합이 K가 된다는 뜻이다.

 2) 만약 있다면, sum[i]-K에 대응하는 value값을 정답에 더해준다.

 3) map에 sum[i]를 넣는다.

 

배운 점

누적 합에 대한 감각이 부족했나보다. prefix sum을 이용하는 누적합 관련 문제를 좀 볼 필요가 있어보인다.

 

코드

import java.io.*;
import java.util.*;
 
public class Main {
 
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken());
        st = new StringTokenizer(br.readLine());
        int[] sum = new int[N];
        sum[0] = Integer.parseInt(st.nextToken());
        for (int i = 1; i < N; i++) {
            sum[i] = sum[i-1] + Integer.parseInt(st.nextToken());
        }
        
        long ans = 0;
        HashMap<Integer, Long> hMap = new HashMap<>();
        // 자기 자신이 K인 경우도 감안하기 위해 넣어준다.
        hMap.put(0, 1L);
        // i = m에서, hMap에는 sum[0], ..., sum[m-1]이 담겨 있음.
           // 따라서, sum[m] - K이 존재한다면
        // 적당한 l<m이 존재하여 sum[m] - K = sum[l], 즉
        // K = sum[m] - sum[l]이라는 뜻이다.
        for (int i = 0; i < N; i++) {
            if (hMap.containsKey(sum[i] - K)) {
                ans += hMap.get(sum[i] - K);
            }
            
            if (hMap.containsKey(sum[i])) {
                hMap.put(sum[i], hMap.get(sum[i]) + 1);
            } else {
                hMap.put(sum[i], 1L);
            }
        }
        System.out.println(ans);
    }
}