[Algorithm] SWEA : 10965. 제곱수 만들기 by Python

엄희관·2020년 12월 28일

Algorithm

목록 보기

43/128

📌문제 설명

어떤 자연수 A가 주어진다. 여기에 자연수 B를 곱한 결과가 거듭제곱수가 되는 최소의 B를 구하는 프로그램을 작성하라.

입력

첫 번째 줄에 테스트 케이스의 수 T 가 주어진다.

각 테스트 케이스의 첫 번째 줄에는 하나의 자연수 A(1≤A≤10**7) 가 주어진다.

💡 문제 풀이

시간초과로 많은 어려움을 겪었던 문제...

1차시도 - 시간초과

문제의 난이도가 높지 않아 단순하게 2부터 A까지 나누었을 때 나눈 횟수가 홀수이면 곱해주는 방식으로 정답을 도출하려 하였다.

T = int(input())
for tc in range(T):
    A = int(input())
    res = 1
    # items = dict()
    for i in range(2, A+1):
        if not A%i:
            cnt = 1
            A //= i
            while True:
                if not A%i:
                    cnt += 1
                    A //= i
                else:
                    break
            if cnt % 2:
                res *= i
            if A == 1:
                break
    print('#{} {}'.format(tc+1, res))

자연수 A의 범위가 10**7 까지이다 보니 큰 숫자의 경우 위의 방법은 시간초과가 발생하여 오답!
※ 평균적으로 100000개의 테스트 케이스 중 190개가 통과하였다...

2차시도 - 시간초과

혹시나해서 나누는 작업을 하지 않고 단순히 숫자를 증가시키면서 제곱수면 출력하도록 아래와 같이 코드를 작성해보았다.

import math

T = int(input())
for tc in range(T):
    A = int(input())
    for i in range(1, A+1):
        if A * i == int(math.sqrt(A * i)) ** 2:
            print('#{} {}'.format(tc+1, i))
            break

예상대로 시간초과가 발생하였고 첫번째 방법보다 훨씬 적은 60개 가량의 테스트 케이스가 통과되었다.

3차시도 - 해결

모든 테스트 케이스에서 1부터 A까지의 범위의 반복문을 돌리는 것이 시간초과의 원인이라고 생각하였다.

따라서, 소수에 대해서만 A가 나누어떨어지는지 확인하려고 했다.
ex) 2로 나눈 이후에는 4, 6, 8.. 과 같은 짝수는 나누어질 수가 없다.

모든 테스트 케이스마다 반복해서 소수를 구하는 것은 비효율적이므로 가장 상단에 2부터 3162( = int(10000000 ** 0.5))까지의 소수를 담은 리스트를 만든다.

prime : 2부터 3162까지의 소수를 담은 리스트

A 값이 제곱수이면 더 이상 진행할 필요가 없지만 제곱수가 아니라면 prime에 있는 소수값들로 하나씩 A값이 나누어지는지 확인한다.
→ 제곱수를 만들기 위한 값을 파악

나누어떨어지는 경우 A가 소수 p에 대해서 몇 번을 나눌 수 있는지 개수를 파악한다.

cnt : 특정 소수가 몇번 나누어떨어지는 확인하기 위한 변수
res : A에 곱한 결과가 거듭제곱수가 되는 최소의 자연수(자연수 B)

이후 cnt의 값이 짝수면 제곱근이 될 수 있지만 홀수면 해당 소수를 곱해줘야 제곱근이 되므로 res에 해당 소수(p)를 곱해준다

시간을 줄이기 위해서 A가 더 나누어 떨어질 수 없는 다음의 두 가지 경우는 break해준다.
1. A == 1
2. p > A

prime에 있는 모든 소수로 나누었음에도 나누어떨어지지 않았다면 해당 값은 prime의 최대값보다 더 큰 소수값이므로 res에 곱해준다.

코드는 아래와 같다

prime = [2]
for i in range(3, int(10000000 ** (0.5)), 2):
    for p in prime:
        if not i % p: break
    else:
        prime.append(i)
answer = []
T = int(input())
for tc in range(T):
    A = int(input())
    res = 1
    if A**0.5 != int(A**0.5):
        for p in prime:
            cnt = 0
            while not A % p:
                A //= p
                cnt += 1
            if cnt % 2:
                res *= p
            if A == 1 or p > A:
                break
        if A > 1:
            res *= A
    answer.append('#{} {}'.format(tc+1, res))
for ans in answer:
    print(ans)