[백준] 2015번 - 수들의 합4 (java)

Baekjoon 2015 - 수들의 합4 (클릭 시 이동)

문제

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인 것이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

시간 제한	메모리 제한	제출	정답	맞은 사람	정답 비율
2 초	128 MB	1901	599	445	33.434%

입력

첫째 줄에 정수 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. 수들의 누적합을 배열에 저장한다.
2. 누적합의 수와 갯수를 저장할 map을 선언한다.
3. 누적합이 K와 같다면 정답에 더해준다.
4. 처음부텨 현재까지 더한 수가 A라고 했을 때 누적합이 A-K인 수가 있으면 그 갯수만큼 정답에 더해준다.

과정

map에 저장할 Key와 value 모두 Long으로 선언한다.

Key는 누적합, value는 갯수인데

만약 N개 모두 값이 10,000이라고 하면 N의 최대값은 200,000 이므로 누적합의 최댓값은 20억이므로 Key는 Integer타입으로 설정하고

200,000개 모두 값이 0이고 비교할 값(K)이 0이라면 count의 갯수는 모든 부분집합의 갯수N*(N+1) / 2개가 되기 때문에 Value는 Long으로 설정 해 준다.

Map<Integer, Long> map = new HashMap<>();
for (int i = 1; i <= N; i++) {
    // 부분합 구하기
    data[i] = Integer.parseInt(st.nextToken()) + data[i - 1];

    // 부분합이 K라면 결괏값에 반영
    if (data[i] == K) {
        answer++;
    }

    // 누적합이 A이고 A-K가 B라고 했을 때
    // 누적된 값이 B인 경우가 있다면 처음부터 B가 되는 위치까지 제외시킨다면 값은 K가 될 것이다.
    // 그 갯수만큼 결과에 반영해준다.
    if (map.containsKey(data[i] - K))
        answer += map.get(data[i] - K);

    // 누적된 수가 존재하면 value에 1을 더해주고 없으면 만들어준다.
    if (!map.containsKey(data[i]))
        map.put(data[i], 1L);
    else
        map.put(data[i], map.get(data[i]) + 1);
}

최종 소스코드

import java.io.*;
import java.util.*;

public class BJ_G5_2015_수들의_합4 {
    static int N, K;
    static int data[];
    static Map<Integer, Long> map = new HashMap<>();
    static StringTokenizer st;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();

        st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());
        data = new int[N + 1];

        long answer = 0;
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) {
            data[i] = Integer.parseInt(st.nextToken()) + data[i - 1];
            if (data[i] == K) {
                answer++;
            }
            if (map.containsKey(data[i] - K))
                answer += map.get(data[i] - K);
            if (!map.containsKey(data[i]))
                map.put(data[i], 1L);
            else
                map.put(data[i], map.get(data[i]) + 1);
        }

        System.out.println(answer);
    }
}