SWEA 13736. 사탕 분배

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AX8BB5d6T7gDFARO&categoryId=AX8BB5d6T7gDFARO&categoryType=CODE&problemTitle=13736&orderBy=FIRST_REG_DATETIME&selectCodeLang=ALL&select-1=&pageSize=10&pageIndex=1&&&&&&&&&

SW Expert Academy

 

전체 코드

def power(K, sum):
    if K == 1:
        return 2

    ret = power(K // 2, sum)
    ret = (ret ** 2) % sum

    if K % 2 == 1:
        ret = (ret * 2) % sum

    return ret


def solution(A, B, k):
    sum = A + B
    kp = power(k, sum)

    A = (A * kp) % sum

    return min(A, sum - A)


T = int(input())
for testcase in range(1, T + 1):
    A, B, K = map(int, input().split())
    result = solution(A, B, K)
    print(f'#{testcase} {result}')

 

풀이

해당 문제는 성능을 위해서 반드시 일반화가 필요하다. 먼저 사탕의 총 개수는 변하지 않는다. A+B=sum이라고 했을 때 사탕을 분배한 후의 개수를 A', B'라고 했을 때 A'+B'=sum이다. 

A<B일 때 A'=2A이고 A>B일 때 A'=A-B이다. 여기서 B=sum-A이므로 A'=A-(sum-A)=2A-sum이다.

A>B, A<B 모두를 만족시키는 일반화된 방법이 있을까? 결과에 모듈러 연산(%)을 진행하면 된다.

[A<B의 경우]

2A < (A+B)이므로 2A % (A+B)는 항상 2A이다. 따라서 모듈러 연산에 영향을 받지 않는다.

A>B의 경우 

모듈러 연산 분배 법칙에 의해 다음과 같이 진행된다. A<B, A>B는 모두 A'=2A % sum 이다.

A'에서 분배를 진행하면 다음과 같이 진행된다.

2^kA % sum을 구하면 k번 교환했을 때 A의 값을 구할 수 있다. B는 B=sum-A이므로 A를 통해 값을 구할 수 있다.

A를 구하는 것은 power함수를 통해 구한다.

def power(K, sum):
    if K == 1:
        return 2

    ret = power(K // 2, sum)
    ret = (ret ** 2) % sum

    if K % 2 == 1:
        ret = (ret * 2) % sum

    return ret

왜 2^k*A % sum으로 바로 값을 구하지 않을까? 이유는 k의 값이 클수록 값이 엄청나게 커지기 때문이다. k = 2000000일 경우 2^2000000 값은 엄청나게 커지기 때문에 바로 구하면 시간이 오래걸린다. 따라서 백트래킹된 값에 "% sum" 모듈러 연산을 진행하여 값을 계산한다. 

K // 2를 재귀함수 파라미터로 전달하여 log K만큼 연산이 진행되도록하며 이를 통해 연산량을 크게 줄일 수 있다. 이 후 반환된 값을 제곱하여 값을 반환한다. 만약 현재 K가 홀수인 경우 K // 2 연산 시 값이 차수가 1 줄어드므로 2를 곱한 후 % sum 연산을 진행한다.

 

결과에 A를 곱한 후 % sum 모듈러 연산을 진행하면 2^kA % sum 값을 구할 수 있다.

def solution(A, B, k):
    sum = A + B
    kp = power(k, sum)

    A = (A * kp) % sum

    return min(A, sum - A)

 

회고

A'=2A % sum으로 일반화하는 작업이 무엇보다 어려웠다. 연산을 반복하는 것이 아닌 일반화된 방정식을 유도하는 것인데 완전히 처음 접해보는 문제라 발상이 떠오르지 않았다.

두번째로 모듈러 연산의 일반화된 분배법칙을 모르는 것이 크다. 사칙연산과 다르게 나머지 연산은 분배법칙이 다르다. 이 때문에 2^k * A % sum을 유도하는 것을 이해하기가 많이 어려웠다.

모듈러 연산에 대한 일반화된 공식에 대해 자료조사가 필요할 듯하다.