[PS] BOJ 29719 / 브실이의 불침번 근무

[PS] BOJ 29719 / 브실이의 불침번 근무
Thumbnail: Photo by Aleksey Kashmar (Unsplash)
문제 링크: https://www.acmicpc.net/problem/29719

큰 수의 제곱 연산을 빠르게 하려면 어떻게 해야 할까요?

풀이

문제 이해하기

$M$명의 군인들이 $N$일동안 불침번 근무를 서야 합니다. 하루에는 1명의 군인만 불침번을 서며, 한 명이 여러 번 근무를 설 수 있습니다.

이는 $M$개의 원소 중에 중복을 허용한 채로 $N$개를 뽑는 문제이며, 전체 경우의 수는 $M^N$이 됩니다.

이 중 브실이가 포함되는 경우의 수를 구해야 합니다.

브실이가 포함되는 경우의 수는 전체 경우의 수에서 브실이가 포함되지 않는 경우의 수를 뺀 것과 같습니다. 브실이가 포함되지 않는 경우의 수는 $M-1$명의 원소 중 $N$개를 뽑는 것이므로 $(M-1)^N$이 됩니다.

따라서, 브실이가 포함되는 경우의 수는 $M$$N$ - $(M-1)$$N$이 됩니다.

빠른 모듈로-제곱 연산

$M$과 $N$이 최대 106으로 주어지므로, 제곱 연산자를 바로 사용할 경우 시간 초과를 받게 됩니다. 따라서, 제곱 연산을 빠르게 진행하기 위해 별도의 함수를 구현했습니다. baseexp만큼 제곱한다고 할 때, 이를 base의 거듭제곱의 곱으로 쪼개어 계산하는 방식입니다.

지수인 exp를 2진수로 표기할 때, $i$번째 자리 비트는 base를 $i$번 거듭제곱한 (2$i$만큼 제곱한) 값을 나타내며, 해당 비트가 1이면 이 값을 결과값에 곱해주고, 아니라면 넘어갑니다. 항상 곱한 뒤에 모듈로 연산을 진행해 결과값을 작게 유지해줍니다.

💡 모듈로 연산의 분배 법칙으로,
$A \times B \mod C = (A \mod C \times B \mod C) \mod C$입니다.
def mod_pow(base, exp):
    # base %= MOD # base는 항상 MOD보다 작은 수로 주어지므로 불필요하다.
    ans = 1
    while exp > 0: # 지수를 2진수로 나타냈을 때의 각 자릿수를 낮은 비트부터 보며 반복한다.
        if exp & 1: # 현재 비트가 1이면, 현재의 거듭제곱 값을 결과값에 곱해준다.
            ans = ans * base % MOD
        base = base * base % MOD # 다음 비트 자리수에 맞게 거듭제곱을 진행한다.
        exp >>= 1 # 다음 비트로 이동한다.
    return ans

빠른 모듈로-제곱 연산 구현

전체 코드

MOD = 1_000_000_007
input = open(0).readline

def mod_pow(base, exp):
    # base %= MOD # base는 항상 MOD보다 작은 수로 주어지므로 불필요하다.
    ans = 1
    while exp > 0: # 지수를 2진수로 나타냈을 때의 각 자릿수를 낮은 비트부터 보며 반복한다.
        if exp & 1: # 현재 비트가 1이면, 현재의 거듭제곱 값을 결과값에 곱해준다.
            ans = ans * base % MOD
        base = base * base % MOD # 다음 비트 자리수에 맞게 거듭제곱을 진행한다.
        exp >>= 1 # 다음 비트로 이동한다.
    return ans

N, M = map(int, input().split())
print((mod_pow(M, N) - mod_pow(M - 1, N)) % MOD)

solution.py