2015. 수들의 합 4

플랫폼

BOJ

분류

답 본 문제

알고리즘

누적합

해시 테이블

중요도

⭐️⭐️

티어

Gold

링크

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

다시 풀어도 틀릴 것 같은 문제

Created At

2024/05/13 13:12

Updated At

2024/05/13 13:28

Writer

상태

완료

Problem

문제

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

Plain Text
복사

예제 출력 1

Plain Text
복사

예제 입력 2

Plain Text
복사

예제 출력 2

Plain Text
복사

Code

#1. 시간초과 코드

Idea

누적합 배열 채우기

이중 for문으로 순회하면서 구간합 구하기

Code

# 시간초과 코드

import sys

read = sys.stdin.readline

n, k = map(int, read().split())

nums = list(map(int, read().split()))

psum = [0] * (n + 1)

cnt = 0

for i in range(n):
    psum[i + 1] = psum[i] + nums[i]

for i in range(n):
    for j in range(i, n):
        if psum[j + 1] - psum[i] == k:
            cnt += 1

print(cnt)
Python
복사

Commentary

N 이 200,000 이므로 O(N^2) 로직은 시간초과가 난다

#2. 정답코드

Idea

구간합을 구하는 공식을 그대로 사용하지 않고 이항을 통해서 우변에 과거 누적합만을 둔다

(아래 손필기 참고)

Code

import sys

read = sys.stdin.readline

n, k = map(int, read().split())

nums = list(map(int, read().split()))

psums = [0] * (n + 1)

psum_cnt = {}

ans = 0

# 누적합 생성
for i in range(n):
    psums[i + 1] = psums[i] + nums[i]

# k = psum[j+1] - psum[i]
# 누적합 배열을 순회하면서 psum - k 가 psum_cnt 에 있는지 검사하고
# 있으면 (psum-k) 의 개수만큼 더해준다
# 그 후 psum_cnt 의 psum 값을 1 증가시킨다 (개수 1 증가)
for psum in psums:
    if psum - k in psum_cnt:
        ans += psum_cnt[psum - k]

    psum_cnt[psum] = psum_cnt.get(psum, 0) + 1

print(ans)
Python
복사

Solution

Commentary

•

구간합을 2중 for문이 아닌 for 문 하나로 구하기 위한 아이디어를 생각하는게 핵심이였음

◦

현재누적합(psum) - k 가 과거누적합 딕셔너리(prefix_cnt) 에 있는지 확인하는 아이디어

•

n,k 가 매우 크기 때문에 단순리스트가 아니라 딕셔너리를 사용해서 값 조회 시간을 O(1) 로 단축

•

아이디어 발상이 매우매우 어려웠고 풀이를 처음에 보고도 이해하기까지도 오래걸렸다