[프로그래머스] 기사단원의 무기 Python 풀이(제곱근 이용)

문제링크: https://school.programmers.co.kr/learn/courses/30/lessons/136798

프로그래머스

 

예제 입력: 

5 3 2

10 3 2

 

예제 출력:

10

21

풀이

프로그래머스 레벨 1 문제라 그런지 문제 자체는 간단하다. 

number 범위에 해당하는 각 숫자의 약수를 구하고, 그 약수의 개수가 제한 수치를 넘는지 확인하여 적절한 무기의 공격력을 결정한 뒤 필요한 철의 무게를 누적하면 되는 문제이다.

 

처음 작성한 코드:

def solution(number, limit, power):
    answer = 0
    divisors = [[] for _ in range(number+1)]

    for i in range(1, number+1):
        for j in range(1, number+1):
            if i % j == 0:
                divisors[i].append(j)
    
    res = []
    for i in divisors:
        res.append(len(i))
            
    for i in res:
        if i > limit:
            answer += power
        else:
            answer += i
            
    return answer

 

하지만... 위의 코드로는 어마무시한 시간초과가 발생한다

 

고민하던중, 질문하기에 작성된 글을 보고 제곱근을 이용해야 한다는 힌트를 얻었다.

왜 제곱근을 이용하냐면, 어떤 수의 약수는 대칭적이라는 특징이 있기 때문이다.

 

예를 들어 36의 약수를 생각해보자.

1, 2, 3, 4, 6, 9, 12, 18, 36

위 약수들을 살펴보면, 1과 36, 2와 18, 3과 12, 4와 9는 서로 곱하면 36이 되는 쌍을 이루고 있다. 그리고 이런 쌍들이 모여서 36의 약수를 구성하게 된다.

여기서 특징적인 것은, 이 쌍들 중에서 가장 가운데 있는 쌍이 제곱근에 가장 가깝다는 것이다. 그래서 약수를 찾을 때는 제곱근까지만 찾아보면, 그 이상은(처음코드) 이미 알고 있는 약수의 쌍을 반복해서 찾게 되므로 비효율적이다.

즉, 제곱근까지만 반복문을 돌면서 약수를 찾는것이 문제를 해결할 수 있는 방법이다. 그리고 찾은 약수를 이용해서 약수의 쌍을 구한다.

만약 n이 완전제곱수라면 (즉, 제곱근이 정수라면) 약수의 쌍이 하나 더 있는 것이므로 count를 1 증가시키고, 그렇지 않다면 약수의 쌍이 두 개 더 있는 것이므로 count를 2 증가시킨다.

 

정답 코드:

def solution(number, limit, power):
    # 각 숫자의 약수의 개수를 계산하는 함수
    def divisor_count(n):
        count = 0
        for i in range(1, int(n**0.5) + 1):
            if n % i == 0:
                # 제곱근이 정수인 경우
                if i == (n // i):
                    count += 1
                # 제곱근이 정수가 아닌 경우
                else:
                    count += 2
        return count

    total_iron = 0
    for i in range(1, number + 1):
        weapon_power = divisor_count(i)
        if weapon_power > limit:
            total_iron += power
        else:
            total_iron += weapon_power
            
    return total_iron

위의 방식으로 약수를 개산하면 O(√n)의 시간복잡도를 가지게 된다. 즉, number가 최대 100,000까지 주어질 수 있기 때문에 문제의 제한 사항 내에서 실행할 수 있게된다.