파이썬으로 a, b, c로 나누어 떨어지는 수열의 n번째 항 구하기

문제 개요

네 개의 숫자 n, a, b, c가 주어졌을 때, a, b 또는 c 중 하나로 나누어 떨어지는 수들을 오름차순으로 정렬한 수열의 n번째 항을 찾는 것이 목표입니다.

예를 들어 입력이 n = 8, a = 3, b = 7, c = 9라면, 조건을 만족하는 수들은 [3, 6, 7, 9, 12, 14, 15, 18] 순서로 나열되며, 이때 정답은 18입니다.

풀이 접근 방법

수열을 처음부터 하나씩 생성하는 것은 비효율적이므로, 포함-배제 원리(Inclusion-Exclusion Principle)와 이진 탐색(Binary Search)을 결합하여 해결합니다.

핵심 아이디어

임의의 값 x에 대해 'x 이하의 수 중 a, b, c 중 하나라도 나누어 떨어지는 수의 개수'는 몫 연산만으로 빠르게 계산할 수 있습니다.
x//a + x//b + x//c에서 두 번씩 중복되어 세어진 부분(x//lcm(a,b) 등)을 빼고, 세 번 빠진 부분(x//lcm(a,b,c))을 다시 더하면 정확한 개수를 얻습니다.
이 개수가 n과 일치하는 지점을 이진 탐색으로 찾습니다.

단계별 알고리즘

a, b, c 중 최솟값이 1이라면 모든 자연수가 조건을 만족하므로 n을 그대로 반환합니다.
ab = lcm(a, b), bc = lcm(b, c), ca = lcm(a, c)를 계산합니다.
abc = lcm(ab, c)를 계산합니다.
탐색 범위를 left = 1, right = 10^9로 설정합니다.
left <= right인 동안 다음을 반복합니다.
- mid = (left + right) // 2
- na = mid // a, nb = mid // b, nc = mid // c
- nab = mid // ab, nbc = mid // bc, nca = mid // ca
- nabc = mid // abc
- numterms = na + nb + nc - nab - nbc - nca + nabc (포함-배제 원리 적용)
- numterms > n이면 right = mid - 1
- numterms < n이면 left = mid + 1
- numterms == n이면 mid - min(mid % a, mid % b, mid % c)를 반환합니다. mid가 조건을 만족하는 수가 아닐 수 있으므로, mid 이하에서 조건을 만족하는 가장 큰 수로 보정하는 과정입니다.
반복문이 종료되면 -1을 반환합니다.

여기서 lcm(최소공배수)은 gcd(최대공약수)를 이용해 (a * b) // gcd(a, b)로 구할 수 있습니다. 탐색 상한인 10^9은 문제의 제약 조건에서 충분히 큰 값으로 설정한 것입니다.

파이썬 구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

import math

def lcm(a, b):
    return (a * b) // math.gcd(a, b)

class Solution:
    def solve(self, n, a, b, c):
        if min(a, b, c) == 1:
            return n
        ab, bc, ca = lcm(a, b), lcm(b, c), lcm(a, c)
        abc = lcm(ab, c)
        left, right = 1, 10 ** 9
        while left <= right:
            mid = (left + right) // 2
            na = mid // a
            nb = mid // b
            nc = mid // c
            nab = mid // ab
            nbc = mid // bc
            nca = mid // ca
            nabc = mid // abc
            numterms = na + nb + nc - nab - nbc - nca + nabc
            if numterms > n:
                right = mid - 1
            elif numterms < n:
                left = mid + 1
            else:
                return mid - min(mid % a, mid % b, mid % c)
        return -1

ob = Solution()
n = 8
a = 3
b = 7
c = 9
print(ob.solve(n, a, b, c))

입력

8, 3, 7, 9

출력

시간 복잡도

매 반복마다 상수 개의 나눗셈 연산만 수행하므로, 전체 시간 복잡도는 탐색 범위에 대한 로그 시간인 O(log 10^9)로 매우 효율적입니다. 이는 수열을 직접 순회하는 O(n) 방식보다 훨씬 빠릅니다.