[백준] 2015번 : 수들의 합4 - 자바(JAVA)

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

2015번: 수들의 합 4

 

문제 해석

A [1], A [2],..... A [N]의 N개의 정수가 저장되어 있는 배열이 있습니다.

이 배열 A의 부분합이란 1 <= i <= j <= N인 정수 i와 j에 대해 A [i]와 A [j]까지의 합을 말합니다.

 

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

 

정수 N은 최대 200,000입니다. (N Log(N) 시간복잡도 이하로 해결해야 함)

K는 20억이하의 수입니다. (21억을 넘지 않기 때문에 int형으로 가능 , long 타입 고려 안 해도 됨)

주어지는 정수의 절댓값은 10,000을 넘지 않습니다.

하지만 출력해야 할 정답은 N*(N+1)/2는 int형 overflow가 발생할 수 있기 때문에 long타입을 사용해야 합니다. (N = 20만)

문제 풀이 전 설계

우선 i와 j의 순서가 중요하기 때문에 정렬은 하면 안될 것 같습니다.

 

(i , j)의 나올 수 있는 순서 쌍은 다음과 같습니다

i=1 일때 j= N-1개

i=2 일때 j = N-2개

i=3 일때 j = N-3개

 

하지만 이렇게 답을 구하면 O(N^2)입니다.

 

정렬이 되지 않기 때문에 단순한 투 포인터로 해결할 수 없을 것 같습니다.

 

이분 탐색을 활용한 파라미터 서치로 답을 찾아가는 것에서도 결국 부분합의 개수를 찾지 않고 답을 찾아가는 방법을 찾아야 하는데 생각이 나지 않습니다.

 

세그먼트 트리도 생각은 나지만 세그먼트 트리는 부분합을 저장해놓고 여러 번의 쿼리가 들어올 때 유의미하게 사용할 수 있을 것 같습니다.

 

문제 풀이

문제에서 주는 알고리즘 분류 힌트는 4가지입니다.

1. 자료 구조

2. 누적 합

3. 해시를 사용한 집합과 맵

4. 트리를 사용한 집합과 맵

 

질문검색까지 들어가 보니 다음과 같은 힌트가 있습니다.

Map으로 각 누적합에 대한 개수를 저장하세요.

 

 

우선 누적합에 대해 알고 있어야 합니다.

i~j구간의 구간합을 구하기 위해서는 누적합의 방법으로 구할 수 있습니다.

 

예를 들어 arr배열에 대하여 index 3~ index 5인 구간의 구간합을 구한다고 가정하겠습니다.

 

이때 arr[0~5]까지의 합 - arr [0~2]까지의 합을 하게 되면  arr [3~5]까지의 합이 나오게 됩니다.

 

i~j의 구간합이 k와 같기 위함을 공식으로 나타내면 다음과 같습니다.

sum[i] - sum [j-1] = k

 

여기까지 봤을때 누적합으로 부분합을 구하는 것이 이해되지 않는다면 아래의 출처에서 누적합을 보고 오거나 개인적으로 찾아보시면 좋을 것 같습니다.

 

이 문제에서는 단순히 구간합을 구하는것이 아니라 구간합이 K인 것의 개수를 구해야 합니다.

 

단순히 O(N^2)으로 반복문을 통해 구현하면 시간초과가 예상됩니다 N = 20만

 //다음 코드와 같이 위의 누적합을 통해 구간합을 O(N^2)을 통해 개수를 구하면 시간초과

for(int i=1; i<=N; i++){
	for(int j=i; j<=N; j++){
		if(prefixSum[i] - prefixSum[j] == k){
			count++;
		}
	}
}

 

다음 아이디어를 활용해야 합니다.

현재 누적합이 A이고 A-K = B라고 하였을 때
누적된 값이 B인 경우가 있다면 A - B = (A) - (A-K) => K

 다음 예제를 예시로 들어보겠습니다

4 0

2 -2 2 -2

 

prefixSum 배열은 다음과 같이 구성될 것입니다.

prefixSum[0] = 0;

prefixSum [1] = 2;

prefixSum [2] = 0;

prefixSum [3] = 2;

prefixSum [4] = 0;

 

이제 prefixSum배열을 토대로 구간합이 (K = 0)인 개수를 찾아보겠습니다.

prefixSum [1] = 2로 0이 아닙니다.

현재 누적합 A = 2입니다.

현재 누적합이 2였던 경우가 존재하지 않기 때문에 카운트되지 않습니다.

map에 누적합이 2였던적이 존재했다고 기록합니다. map.put(2, 1);

 

prefixSum[2] = 0입니다.

즉, result++ 연산이 일어납니다.

현재 누적합이 0이었던 적이 존재하지 않기 때문에 추가로 카운트되지 않습니다.

map에 누적합이 0이었던 적이 존재한다고 기록합니다. map.put(0,1);

 

prefixSum [3] = 2입니다.

K값이 아니기 때문에 result++ 연산이 일어나지 않습니다.

현재 누적합 A =2이며 누적합이 2였던 적이 존재하기 때문에 추가로 카운트됩니다.

result += map.get(prefixSum [i] - K);

이후에 누적합이 2였던 적을 추가로 +1 기록합니다.

map.put( 2, map.get(prefixSum [i]) + 1); 

 

이 과정을 반복하면 O(N^2)만큼 탐색하면서 검사할 필요가 없이 이전에 나왔던 누적 합의 개수들을 Map에 기록하기 때문에 A-B = K 조건을 만족하는 누적합이 존재하는 만큼 result ++을 해줄 수 있게 됩니다.

 

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;

public class Main_2015번_수들의합4 {
    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 = 0;
        int[] prefixSum = new int[N + 1];
        for (int i = 1; i <= N; i++) {
            sum += Integer.parseInt(st.nextToken());
            prefixSum[i] = sum;
        }

        //입력 끝 + 누적합을 구함

        //다음 코드와 같이 위의 누적합을 통해 구간합을 O(N^2)을 통해 개수를 구하면 시간초과
        /*
            for(int i=1; i<=N; i++){
                for(int j=i; j<=N; j++){
                    if(prefixSum[i] - prefixSum[j] == k){
                        count++;
                    }
                }
            }
         */

        long result = 0L;
        Map<Integer, Long> map = new HashMap<>();

        for (int i = 1; i <= N; i++) {
            //만약 누적합이 K라면 결괏값에 반영
            if (prefixSum[i] == K) {
                result++;
            }

            //누적합이 A이고 A-K = B라고 하였을 때
            //누적된 값이 B인 경우가 있다면 A - B = (A) - (A-K) => K
            //따라서 그 개수만큼 결과값에 반영
            result += map.getOrDefault(prefixSum[i] - K, 0L);
            map.put(prefixSum[i], map.getOrDefault(prefixSum[i], 0L) + 1);
        }

        System.out.println(result);


    }
}

 

출처

https://jow1025.tistory.com/47

누적합(prefix sum)