백준 2015 수들의 합 4 Kotlin (누적 합)

문제 출처 : 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[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. 알고리즘 분류를 확인한다. 혹은 질문 게시판을 확인한다.

2. 다른 사람의 풀이에서 사용한 알고리즘 혹은 자료구조를 확인한다.

3. 다른 사람의 풀이를 확인한다.

누적 합(부분 합)과 투 포인터를 공부하기 위해 푸려고 한 문제이기에 알고리즘 분류는 알고 있었고, 다른 사람의 풀이에서 map을 사용한 것을 확인했다. 그러고 다시 풀이를 생각하는데 도저히 생각이 나지 않아서 결국 풀이까지 봤는데도 이해하는 데 꽤 시간이 걸렸다.

우선 누적 합에서 (i~j)의 구간 합이 k와 같으려면 다음과 같은 식을 만족해야 한다.

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

주어진 입력으로 살펴보자. 

4 0

2 -2 2 -2 

preSum = {0, 2, 0, 2, 0}

//실제 사용할 누적 합은 preSum[1]부터이며, preSum[0]은 누적 합을 만들 때 편의를 위해 넣어둔다.

1~1까지의 구간 합은 2이다.

1~2까지의 구간 합은 0이므로 count를 1 늘려준다.

1~3까지의 구간 합은 2이다.

1~4까지의 구간 합은 0이므로 count를 1 늘려준다.

2~2까지의 구간 합은 -2이다.

2~3까지의 구간 합은 0이므로 count를 1 늘려준다.

2~4까지의 구간 합은 -2이다.

3~3까지의 구간 합은 2이다.

3~4까지의 구간 합은 0이므로 count를 1 늘려준다.

이런 식으로 풀면 N*(N+1)/2로 O(N^2)의 복잡도를 갖게 되어 시간 초과가 뜰 것이다.

 

여기서 시간 복잡도를 줄이기 위해 map을 사용한다.

풀이는 다음과 같다.

1. 1부터 n번째까지의 누적 합을 모두 돌면서 map에 구간 합 num에 대한 개수를 저장한다.

2. 1부터 x번째까지의 누적 합이 k와 같다면 count를 1 증가 시킨다.

3. x번째까지의 구간 합에 k를 뺀 값(== pSum[j-1])의 개수만큼 count를 1 증가시킨다. 

 

x번째까지의 누적 합을 검사할 때, pSum[x] - pSum[j-1] = k의 식을 이용한다.

예를 들어 3번째까지의 누적 합을 검사할 때, 누적 합은 2이다. 1부터 3까지의 누적 합 중에서

k(0)가 되는 구간 합이 있는지 검사해 줘야 하는데, 이를 반복문 대신 map으로 구한다.

위의 구간 합을 구하는 식을 바꾸면 pSum[x] -k = pSum[j-1] 이 된다.

pSum[3]은 2이기 때문에 pSum[3] -pSum[j-1] = 0이 되는 개수는 pSum[j-1]이 2인 구간 합의 개수가 된다.

// 1<= j <=3

즉 pSum[3]-k는 2이기 때문에 pSum[j-1]이 2가 되는 개수만큼 k가 되는 구간 합이 있다는 건데,

1번 과정에서 map에 넣어 놓은 map[2][1] -> 2의 개수가 몇 개인지만 확인하면 된다!

 

 

i =1일 때,

pSum[i]=2

count=0

map[2]=1

 

i =2일 때,

pSum[i] = 0

count =1

map[0] =1

 

i =3일 때,

pSum[i] = 2

count = 2 //map[2]의 개수가 1이었기 때문에 1 증가

map[2]=2

 

 

i=4일 때,

pSum[i] = 0

count = 4 //pSum[i]가 0이므로 1 증가, map[0]의 개수는 1이었기 때문에 1 증가

map[0] =2

 

 

코드

import java.util.StringTokenizer
fun main() = with(System.out.bufferedWriter()){
    val br = System.`in`.bufferedReader()
    var tk = StringTokenizer(br.readLine())
    val (N,K) = List(2){Integer.parseInt(tk.nextToken())}
    //누적합(부분합)
    val pSum = LongArray(N+1)
    var answer=0L
    val map = mutableMapOf<Long,Long>()
    tk = StringTokenizer(br.readLine())

    for(i in 1 .. N){
        pSum[i] = pSum[i-1]+Integer.parseInt(tk.nextToken())
        if(pSum[i]==K.toLong()){
            answer++
        }
        answer += map.getOrDefault(pSum[i]-K,0)
        map.put(pSum[i],map.getOrDefault(pSum[i],0)+1)
    }
    write("$answer")
    close()
}