[ BOJ ][JAVA][2015] 수들의 합 4

www.acmicpc.net/problem/2015

2015번: 수들의 합 4

시간 제한	메모리 제한	제출	정답	맞은 사람	정답 비율
2 초	128 MB	1473	443	337	33.903%

문제

A[1], A[2], ..., A[N]의 N개의 정수가 저장되어 있는 배열이 있다. 이 배열 A의 부분합이란 1 ≤ i ≤ j ≤ N인 정수 i와 j에 대해 A[i]부터 A[j]까지의 합을 말한다.

N과 A[1], A[2], ..., A[N]이 주어졌을 때, 이러한 N×(N+1)/2개의 부분합 중 합이 K인 것이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 정수 N과 K가 주어진다. (1 ≤ N ≤ 200,000, |K| ≤ 2,000,000,000) N과 K 사이에는 빈칸이 하나 있다. 둘째 줄에는 배열 A를 이루는 N개의 정수가 빈 칸을 사이에 두고 A[1], A[2], ..., A[N]의 순서로 주어진다. 주어지는 정수의 절댓값은 10,000을 넘지 않는다.

출력

첫째 줄에 합이 K인 부분합의 개수를 출력한다.

예제 입력 1

4 0
2 -2 2 -2

예제 출력 1

4

코드

/*
    수가 자연수일 때는 투 포인터로 구할 수 있다.
    하지만 이 문제는 음수와 0이 있기 때문에 구간 Map을 이용한 구간 합을 사용해야 한다.
*/
import java.io.*;
import java.util.*;
public class p2015 {
    static int n, k;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());
        int[] arr = new int[n + 1];
        int[] psum = new int[n + 1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
            psum[i] = psum[i - 1] + arr[i];
        }
        // i        :    0    1    2    3
        // psum[i]    :    2    0    2    0
        HashMap<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);
        long ret = 0;
        for(int i = 1; i <= n; i++) {
            ret += map.getOrDefault(psum[i] - k, 0);

            // key에 해당하는 value가 없으면 1, 있으면 value + 1
            map.put(psum[i], map.getOrDefault(psum[i], 0) + 1);
        }

        System.out.println(ret);
    }

}