문제
N개의 정수 수열에서 연속 부분 수열의 합이 K인 것의 개수를 구하라.
입력
첫째 줄에 N과 K가 주어진다 (1 이상 200,000 이하). 둘째 줄에 N개의 정수가 주어진다.
출력
합이 K인 연속 부분 수열의 개수를 출력한다.
예제
| 입력 | 출력 |
|---|---|
5 0 -1 0 1 0 -1 | 5 |
풀이
누적 합과 HashMap으로 합이 K인 구간을 O(N)에 카운트한다.
data[i]를 1~i까지의 누적 합으로 구성한다- 구간
[j+1, i]의 합이 K이면data[i] - data[j] = K, 즉data[j] = data[i] - K이다 - HashMap에 지금까지 등장한 누적 합의 빈도를 저장한다
- 각 i에서
data[i] == K이면 1~i 전체 구간이므로 ans 증가 map.get(data[i] - K)이 있으면 해당 빈도만큼 ans 증가data[i]의 빈도를 map에 갱신한다
핵심 아이디어: 누적 합의 차이가 K인 쌍의 개수를 세는 문제로, HashMap으로 이전 누적 합을 O(1)에 조회하여 O(N)에 해결한다.
코드
package day349;
import java.io.*;
import java.util.*;
public class Day345BOJ2015수들의합4 {
static int N, K;
static int data[];
static Map<Integer, Long> map = new HashMap<>();
static StringTokenizer st;
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
st = new StringTokenizer(br.readLine());
N = Integer.parseInt(st.nextToken());
K = Integer.parseInt(st.nextToken());
data = new int[N + 1];
long ans = 0;
st = new StringTokenizer(br.readLine());
for (int i = 1; i <= N; i++) {
data[i] = Integer.parseInt(st.nextToken()) + data[i - 1];
if (data[i] == K) {
ans++;
}
if (map.containsKey(data[i] - K))
ans += map.get(data[i] - K);
if (!map.containsKey(data[i]))
map.put(data[i], 1L);
else
map.put(data[i], map.get(data[i]) + 1);
}
System.out.println(ans);
}
}복잡도
- 시간: O(N)
- 공간: O(N)