[알고리즘] 고속 거듭제곱

문제

양의 정수 M, e, n, d에 대하여 C = M^e mod n을 구하시오. (mod는 나머지를 구하는 연산이다). 그리고 위의 C에 대하여 C^d mod n을 구하시오.
단, 양의 정수 a, k, n에 대하여 ak mod n을 구하는 방법은 다음을 이용한다.
a^k mod n = a mod n if k = 1
= (a^2/k mod n)^2 mod n if k가 짝수
= (a*(a^k-1 mod n)) mod n if k가 홀수

입력 예시

15 3 7 33 // M e d n

출력 예시

9
15

나의 풀이

처음에 그냥 쉽게 생각했다.

1. 거듭제곱을 구하는 함수 작성(재귀함수를 이용)
2. 문제의 조건을 그대로 mod 함수로 만들어서 mod한 값을 다시 mod 함수에 넣어서 print하기

이렇게 생각하여

def exp(x, n):
    if n == 0:
        return 1
    else:
        return x * exp(x, n - 1)

M, e, d, n = map(int, input().split())

def mod(w, x, y):
    if x == 1:
        return w % y
    elif x % 2 == 0:
        return exp(exp(w, x / 2) % y, 2) % y
    else:
        return w * (exp(w, x - 1) % y) % y

print(mod(M, e, n))
print(mod(mod(M, e, n), d, n))

으로 코드를 짰다. 풀면서 이렇게 쉽게 풀라고 있는 문제는 아니라고 생각했는데, 아니나 다를까 테스트 케이스에서 통과하지 못했다. 절반은 통과하고 나머지 절만은 런타임 에러가 떴는데, 어찌됐든 제대로 된 코드는 아니었다는 생각이 들었다. 그래서 어떻게 할까..를 생각해봤는데, 고속 거듭제곱 알고리즘을 통해 해결할 수 있었다.

M, e, d, n = map(int, input().split())

def mod(a,b,n):    # a^b mod n을 나타낸 함수
    result = 1    # result 변수를 선언하고 값 1을 할당
    while b > 0:    # 지수가 양수인 경우
        if b % 2 != 0:    # 지수가 홀수라면
            result = (result * a) % n    # 기존 result 값에 밑 a값을 곱해주고, 그 값을 n으로 나눈 나머지를 result에 할당
        b = b // 2    # b를 2로 나눈 몫을 새롭게 b로 지정
        a = (a * a) % n    # a를 제곱한 뒤 그 값을 n으로 나눈 나머지를 새롭게 a로 정의

    return result    # b가 0일 경우 while문을 끝내고 result 값을 할당


print(mod(M, e, n))
print(mod(mod(M, e, n), d, n))

고속 거듭제곱은, 분할과 정복을 이용하여 x ** n과 같이 O(n)의 시간 복잡도가 걸려서 여러 수의 거듭 제곱이나, 또는 거듭제곱의 수가 매우 클 때 효율적인 연산이 불가능한 문제를 해결하는 방법이다. (x^a)^b = x^(a*b) 을 이용한 방식이다.

2의 16제곱이 있다고 해 보자. 원래대로라면 2를 16번 곱해주는 형태로 진행이 되겠으나, 지수를 반으로 나누는 방식으로 마지막에 (((2^2)^2)^2)^2 의 형태로 계산할 경우, 4번만에 끝나게 된다.
이 고속 거듭제곱 알고리즘은, O(log2^n)의 시간 복잡도를 갖게 되어 기존의 O(n)보다 훨씬 개선된 알고리즘이다.

이 문제를 고속 거듭제곱을 이용해 계산한 결과 모든 테스트 케이스를 통과할 수 있게 되었다.