정올 5656. 보석 괴물

문제 링크

문제

보석괴물의 몸에는 부터 사이의 숫자가 적혀 있다고 한다.

보석괴물이 사라지기 전에 보석괴물의 몸통 구간에 새겨진 숫자의 합과 몸통 구간의 길이가 일치하는 모든 경우의 수를 정확하게 말한다면 소원을 이루어주고 사라진다고 한다.

몸통 길이를 이라고 할 때 그 범위는 다음과 같다.

의 몸통 구간이 있을 때, 이 몸통 구간의 합과 구간의 길이가 일치하는 모든 경우의 수를 구하여라!

입력

첫 번째 줄에는 배열의 길이 이 입력된다.

두 번째 줄에는 개의 십진수로 구성된 문자열이 주어진다.

여기서 번째 숫자는 의 값과 같다.

출력

각 테스트 케이스에 대해 구간 합과 구간의 길이가 일치하는 경우의 수를 출력한다.

풀이

구간합과 구간길이가 일치하는 부분 구간의 모든 경우를 구하기 위해서는 이전 구간의 정보를 알고 있어야 한다. 누적합 원리를 사용해서 이전의 (0..j) 까지의 구간합과 구간길이 정보를 알고 있다면, 현재 (0..i) 구간에서 (0..j) 구간의 정보를 빼 (i+1..j) 의 구간에 대해 쉽게 구할 수 있다.

한 구간의 끝점에서, 구간합과 구간길이가 같은 부분 구간은 여러 개일 수 있다. N의 최대가 10만인 점을 이용해 최대한 상수 복잡도 안에서 끝내고자 정보를 압축했다.

사용한 아이디어는 다음과 같다.

7
1131113

인덱스를 하나씩 검사할 때마다 (0번부터의 구간길이-0번부터의 구간합) 을 1씩 카운트한다.

idx 0 (길이 1) : 누적합 1 (1-1) -> cnt[0] += 1
idx 1 (길이 2) : 누적합 2 (2-2) -> cnt[0] += 1
idx 2 (길이 3) : 누적합 5 (3-5) -> cnt[-2] += 1
idx 3 (길이 4) : 누적합 6 (4-6) -> cnt[-2] += 1
idx 4 (길이 5) : 누적합 7 (5-7) -> cnt[-2] += 1
idx 5 (길이 6) : 누적합 8 (6-8) -> cnt[-2] += 1
idx 6 (길이 7) : 누적합 11 (7-11) -> cnt[-5] += 1

이 때, (0번부터의 구간길이-0번부터의 구간합) 이 이전에 같았던 경우의 수 = 그 시점부터 부분 구간을 시작하면 구간길이와 구간합이 같아지는 경우 이다.

idx 5 를 탐색 시작할 때에서 예를 들어보자. 
현재까지 쌓인 cnt[0] 은 2, cnt[-2] 는 3이다.

idx 5에서 길이는 6이고, 누적합은 8이기 때문에 (0번부터의 구간길이-0번부터의 구간합) 은 -2 다.

이전에 같았던 경우를 살펴보자.
- idx 2 : 113 까지의 구간이다. 113111에서 해당 구간을 제외하면
  113[111] 남은 구간은 111로, 길이가 3이고 누적합이 3이다.
- idx 3 : 1131 까지의 구간이다. 113111에서 해당 구간을 제외하면
  1131[11] 남은 구간은 11로, 길이가 2이고 누적합이 2이다.
- idx 4 : 11311 까지의 구간이다. 113111에서 해당 구간을 제외하면
  11311[1] 남은 구간은 1로, 길이가 1이고 누적합이 1이다.

따라서, idx 5 를 끝점으로 했을 때 길이==누적합인 부분 구간의 경우는 3개이다.

이 아이디어를 사용해 원리를 적용하면 다음과 같다.

1. for 0..n 에서 현재 인덱스까지의 누적합을 구한다.
2. 현재 인덱스를 끝점으로 한 부분 구간에서 정답에 맞는 경우는 
   cnt[0번부터의 구간길이 - 0번부터의 구간합] 이다.
3. cnt[0번부터의 구간길이 - 0번부터의 구간합] += 1

주의할 점이 있다. 0번부터의 구간길이 - 0번부터의 구간합은 음수가 될 수 있다. 범위를 잘 살펴보면, 해당 값은 최대 100000 ~ 최소 -899999 이 될 수 있다. cnt 배열에 저장할 때, 보정치 90만을 더해서 계산해준다.

• 정답 / PyPy3 / 351ms / 90.4MB
• 정답 / Python 3 / 219ms / 19.1MB

# 5656 : 보석 괴물
import sys
input = sys.stdin.readline

n = int(input().rstrip())
arr = list(map(int, list(input().rstrip())))

cnt = [0 for _ in range(1000001)]
cnt[900000] = 1 # 구간길이가 구간합과 같은 경우도 1 추가해야 함
sums = 0
res = 0

for i in range(n):
    sums += arr[i] # 1. 누적합
    res += cnt[(i+1)-sums + 900000] # 2. 경우 세기
    cnt[(i+1)-sums + 900000] += 1 # 3. 경우 추가

print(res)