[백준/파이썬] 16139번: 인간-컴퓨터 상호작용

문제

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

풀이

• 특정 문자열 $S$, 특정 알파벳 $\alpha$, 문자열의 구간 $[l,r]$
• $S$의 $l$번째 문자부터 $r$번째 문자 사이에 $\alpha$가 몇 번 나타나는가

문제를 읽었을 때 문자열에서 누적합을 이용하면 될 것 같았습니다.
그러나, 어떻게 접근해야 될지 감이 잡히지 않아 슬라이싱과 count()를 이용해 문제를 풀어보았습니다.

import sys
input = sys.stdin.readline

s = input().rstrip()
for _ in range(int(input())):
    a,l,r = input().split()
    l,r = map(int,(l,r))
    tmp = s[l:r+1]
    print(tmp.count(a))

이 코드의 문제점은 매번 문자열에서 [l:r+1] 구간을 잘라내는 작업은 $O(r - l + 1)$의 시간 복잡도가 필요합니다.
최악의 경우 n길이만큼의 구간을 잘라내어 q번 반복하기 때문에, $O(Q × N)$의 시간 복잡도를 갖게 됩니다.

효율적인 코드를 작성하기 위해서 알파벳 소문자의 빈도수를 계산할 배열을 생성했습니다.

cnt = [[0] * (len(s)+1) for _ in range(26)] # 알파벳 소문자 26개에 등장 횟수 기록할 배열
for i in range(len(s)):
    idx = ord(s[i]) - ord('a') # 현재 문자의 알파벳 인덱스
    for j in range(26):
        cnt[j][i+1] = cnt[j][i] # 이전 값들 복사
    cnt[idx][i+1] += 1 # 현재 문자의 등장 횟수 +1

cnt[j][i]는 첫번째 구간부터 i-1번째 구간까지의 알파벳 j가 등장한 횟수입니다.

등장 횟수의 누적합을 배열에 저장하였으니 이를 이용하여 출력해주면 됩니다.

for _ in range(q):
    a,l,r = input().split()
    l,r = map(int,(l,r))
    idx = ord(a) - ord('a') # 입력 받은 a의 인덱스 값
    print(cnt[idx][r+1] - cnt[idx][l]) # a가 l~r구간에 몇 번 등장하는지

코드

# PyPy3
import sys
input = sys.stdin.readline

s = input().rstrip()
q = int(input())
cnt = [[0] * (len(s)+1) for _ in range(26)]

for i in range(len(s)):
    idx = ord(s[i]) - ord('a')
    for j in range(26):
        cnt[j][i+1] = cnt[j][i]
    cnt[idx][i+1] += 1

for _ in range(q):
    a,l,r = input().split()
    l,r = map(int,(l,r))
    idx = ord(a) - ord('a')
    print(cnt[idx][r+1] - cnt[idx][l])