[PS] BOJ 16139 / 인간-컴퓨터 상호작용

문제 링크: https://www.acmicpc.net/problem/16139

Thumbnail: Photo by Thanuj Mathew / Unsplash

'누적 합'에 대한 이해만 있다면 그리 어렵지 않게 풀이를 떠올릴 수 있습니다.

풀이

최대 200,000자의 문자열 \(S\) 에서, 최대 200,000 번의 구간 쿼리를 수행해야 한다.​ 각각의 쿼리에서는 \([L, R]\) (L, R 모두 포함) 의 구간 내의 임의의 알파벳 \(\alpha\)가 몇 개 포함되어 있는지 개수를 세어야 한다.

만약 단순히 구간을 순회한다면, 시간 복잡도는 \(O(NM)\) 형태로 표현할 수 있다.
(N개의 쿼리에 대해, M번씩 반복하며 알파벳 개수를 센다)
하지만 이런 방식으론 금방 시간 초과가 발생하므로, 더 나은 방법을 찾아봐야 한다.

문제를 보면, 초기 문자열 \(S\)의 내용은 변하지 않는다. 즉, \(S\)의 임의의 구간 내의 임의의 알파벳 \(\alpha\)의 개수는 항상 변하지 않는다. 그렇다면, 쿼리 수행 전에 이를 미리 계산해두고 쿼리 시에는 그냥 출력만 하면 되지 않을까?

\(S\)의 길이가 \(N\)이라 할 때, \(S\)를 한 문자씩 누적해가며, \([0, 0]\), \([0, 1]\), \(\cdots\), \([0, N]\) 의 각 구간 안에 모든 종류의 알파벳 소문자가 몇 개 있는지 개수를 세어 두자.

# [각 문자에 대한 누적 개수] 의 누적 합
acc_count = [[0 for _ in range(26)] for _ in range(len(S))]

for i, c in enumerate(S):
    for j in range(26):   # a ~ z 에 대한 누적 개수를 계산한다.
        if i == 0:
            acc_count[i][j] = 1 if ord(c) == ord("a") + j else 0
        else:
            acc_count[i][j] = acc_count[i - 1][j] + (1 if ord(c) - ord("a") == j else 0)

이후, 구간 \([L, R]\)에 문자 \(\alpha\)가 몇개 있는지 세려면, 아래 식으로 계산할 수 있다.

idx = ord(c) - ord("a")
if L > 0:
    print(acc_count[R][idx] - (acc_count[L - 1][idx]))
else: # L이 0인 경우, 그냥 0~R까지의 누적합을 바로 출력하면 된다.
    print(acc_count[R][idx])

코드

from sys import stdin


S = stdin.readline().strip()
Q = int(stdin.readline().strip())
acc_count = [[0 for _ in range(26)] for _ in range(len(S))]     # [각 문자에 대한 누적 개수] 의 누적 합

for i, c in enumerate(S):
    for j in range(26):   # a ~ z 에 대한 누적 개수를 계산한다.
        if i == 0:
            acc_count[i][j] = 1 if ord(c) == ord("a") + j else 0
        else:
            acc_count[i][j] = acc_count[i - 1][j] + (1 if ord(c) - ord("a") == j else 0)

for i in range(Q):
    c, L, R = stdin.readline().strip().split()
    L = int(L)
    R = int(R)
    idx = ord(c) - ord("a")
    if L > 0:
        print(acc_count[R][idx] - (acc_count[L - 1][idx]))
    else:
        print(acc_count[R][idx])