[BOJ] 17435 - 합성함수와 쿼리 (Python)

 

 

아이디어

 

희소 배열을 이용하여 쿼리 수행 시간을 단축한다.

 

 

풀이

 

희소 배열 f의 첫 행에 입력 받은 값들을 저장하고, for문을 돌면서 동적 계획법으로 희소 배열의 행을 추가해 나간다.

희소 배열의 행 크기는 n이 범위를 넘지 않도록 19로 설정했다(2^19 = 524,288 > 500,000).

차례로 n을 1과 AND 비트 연산하여 각 자리값에 따라 희소 배열 f에서 계산된 값을 찾아 나가도록 한다.

 

import sys

input = sys.stdin.readline

m = int(input())
f = [[0]+list(map(int, input().split()))]
for i in range(1, 19):
    f.append([])
    for j in range(m+1):
        f[-1].append(f[-2][f[-2][j]])
Q = int(input())
for _ in range(Q):
    n, x = map(int, input().split())
    for i in range(19):
        if n & 1:
            x = f[i][x]
        n >>= 1
    print(x)

 

 

bit_length 메서드는 해당 수를 2진수로 나타냈을 때의 문자열 기준 길이를 나타낸다.

for i in range(n.bit_length()):
    if n & 1:
        x = f[i][x]
    n >>= 1