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

입력

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

출력

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

 

 

 

누적합을 활용하는 문제다.

정수가 저장된 배열이 입력으로 주어졌을 때, K가 되는 연속된 배열의 부분합이 몇개인지 구하는 문제다.

부분합을 직접 구하려면 O(N^2)의 시간이 드는데, 문제에서 N의 최대가 200000이기 때문에 시간초과가 난다.

즉, 누적합을 활용해야 한다.

 

arr[i] += arr[i - 1] 

위 처럼 자신의 값에 이전의 모든 원소를 저장하여 누적합 해준다.

 

누적합한 배열을 0부터 N-1까지 탐색한다.

만약 누적합 해준 값 arr[i]가 이미 K라면, answer에 1을 더한다.

그리고 현재까지 나온 arr[i] - K의 개수를 answer에 더해주고,

arr[i]의 개수를 +1로 갱신한다.

 

누적합 배열에서 나온 원소와 개수를 저장해주기 위해 HashMap을 사용한다.

key를 원소, value를 개수로 저장해주고,

원소의 개수를 편하게 가져오기 위해 getOrDefault 함수를 사용한다.

 

 

이때 arr[i] - K의 개수를 answer에 더해주는 이유를 이해하기 위해서는 누적합의 아래 특성을 알아야 한다.

 

***

arr[5]는 0~5 인덱스의 부분합이다. (0, 1, 2, 3, 4, 5)

arr[3]은 0~3 인덱스의 부분합이다. (0, 1, 2, 3)

즉, arr[5]와 arr[3]은 0, 1, 2, 3이 겹쳐있으므로

arr[5] - arr[3]은 4~5 인덱스의 부분합이다.

 

 

이 특성을 이해했다면  arr[i] - K 개수를 왜 더해주는지 감이 잡힐 것이다.

arr[i]에서 특정 값(arr[0] ~ arr[i-1])을 빼주면 그것도 이 배열의 부분합이 된다.

즉, arr[i]에서 빼주었을 때 K가 되는 값의 개수를 찾으면

K가 만들어질 수 있는 부분합의 개수가 되는 것이다.

arr[i] - (arr[i] - K) = K가 되기 때문에 arr[i] - K 의 개수를 answer에 더해준다.

 

arr[i]가 이미 K인 경우에도 위 과정을 해주어야 하는데,

arr[0] ~ arr[i-1]에 0이 있다면 K - 0 = K가 되어 K의 부분합을 더 세어줄 수 있다.

 

 

 

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

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());
		
		long[] arr = new long[N];
		st = new StringTokenizer(br.readLine());
		arr[0] = Long.parseLong(st.nextToken());
		
		for(int i=1;i<N;i++) 
			arr[i] = arr[i - 1] + Long.parseLong(st.nextToken());
		
		HashMap<Long, Long> map = new HashMap<>();
		long answer = 0;
		
		for(int i=0;i<N;i++) {
			if(arr[i] == K)
				answer++;
			answer += map.getOrDefault(arr[i] - K, (long)0);
			map.put(arr[i], map.getOrDefault(arr[i], (long)0) + 1);
		}
		
		System.out.println(answer);
	}

}