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

목차

문제 : boj2015

 

 

필요 알고리즘

• 누적 합 (prefix sum), 해시를 사용한 집합과 맵 (hashmap)
  
  • 누적합과 해시를 사용해 풀 수 있는 문제이다.

※ 제 코드에서 왜 main 함수에 로직을 직접 작성하지 않았는지, 왜 Scanner를 쓰지 않고 BufferedReader를 사용했는지 등에 대해서는 '자바로 백준 풀 때의 팁 및 주의점' 글을 참고해주세요. 백준을 자바로 풀어보려고 시작하시는 분이나, 백준에서 자바로 풀 때의 팁을 원하시는 분들도 보시는걸 추천드립니다.

 

 

풀이

   문제에서 A배열에서 1부터 X번까지의 누적합을 f(x)라 해보자. 즉 f(x)는 다음과 같다.

 

  이 때 1 <= i <= j <= N인 정수에 대한 A[i]부터 A[j]까지의 합은 f(j) - f(i-1) 이라고 할 수 있다 (혹시 이해가 안된다면 '누적 합(prefix sum), 2차원 누적합(prefix sum of matrix) with java' 글을 확인해보자.).

 

  그렇다면 이 문제에서 입력을 받을 때, A[1], A[2],... 이 아니라 f(1), f(2), ... 을 미리 저장해둔다고 해보자. 즉 입력받으면서 누적합 배열을 미리 만드는거다. 그럼 '부분합 중 합이 K인 것이 몇 개나 있는지' 라는 이 문제에서 물어보는 것은, '배열을 차례대로 보고 있을 때, 현재 x번째를 보고 있다면 배열에서 이전에 f(x) - k가 몇 번 나왔는가?' 로 바꿔서 풀 수 있게 된다. ( f(x) - 찾는값 = k 이므로, 찾는값 = f(x) - k)

 

  그럼 로직은 다음과 같다.

1. 누적합 배열을 만든다.

2. 누적합 배열을 하나씩 보면서, 현재 x번째를 보고 있다면 arr[x] - K 가 이전까지 몇 번 나왔는지를 답에 더한다.

3. arr[x]가 나왔다고 카운팅한다.

for (int i = 1; i <= n; i++) {
    arr[i] = arr[i-1] + Long.parseLong(st.nextToken());	// [1]
}

long answer = 0l;
Map<Long, Long> cnt = new HashMap<>();
cnt.put(0l, 1l);
for (int i = 1; i <= n; i++) {
    answer += cnt.getOrDefault(arr[i]-k, 0l);	// [2]
    cnt.put(arr[i], cnt.getOrDefault(arr[i], 0l) + 1);	//[3]
}
System.out.println(answer);	// 답 출력

 

 

 

코드 : github

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

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in), 1 << 16);

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

    public void solution() throws Exception {
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        long k = Long.parseLong(st.nextToken());
        long[] arr = new long[n+1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = arr[i-1] + Long.parseLong(st.nextToken());
        }

        long answer = 0l;
        Map<Long, Long> cnt = new HashMap<>();
        cnt.put(0l, 1l);
        for (int i = 1; i <= n; i++) {
            answer += cnt.getOrDefault(arr[i]-k, 0l);
            cnt.put(arr[i], cnt.getOrDefault(arr[i], 0l) + 1);
        }
        System.out.println(answer);
    }
}