BOJ_2015_수들의 합 4 (Java)

[Gold IV] 수들의 합 4 - 2015

문제 링크

성능 요약

메모리: 35308 KB, 시간: 352 ms

분류

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

제출 일자

2024년 11월 25일 22:31:56

문제 설명

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인 부분합의 개수를 출력한다.

코드

/**
 * Author: nowalex322, Kim HyeonJae
 */
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileInputStream;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;

public class Main {
    static BufferedReader br;
    static BufferedWriter bw;
    static StringTokenizer st;
    static int N;
    static long K, prefixSum[];

    public static void main(String[] args) throws Exception {
        new Main().solution();
    }

    public void solution() throws Exception {
        br = new BufferedReader(new InputStreamReader(System.in));
//        br = new BufferedReader(new InputStreamReader(new FileInputStream("src/main/java/BOJ_2015_수들의합4/input.txt")));
        bw = new BufferedWriter(new OutputStreamWriter(System.out));

        st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        K = Long.parseLong(st.nextToken());
        prefixSum = new long[N + 1];

        Map<Long, Long> sumCnt = new HashMap<>();
        long res = 0;

        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) {
            prefixSum[i] = prefixSum[i - 1] + Long.parseLong(st.nextToken());

            if (prefixSum[i] == K) res++;

            res += sumCnt.getOrDefault(prefixSum[i] - K, 0L);
            sumCnt.put(prefixSum[i], sumCnt.getOrDefault(prefixSum[i], 0L) + 1);
        }

        bw.write(String.valueOf(res));
        bw.flush();
        bw.close();
        br.close();
    }
}