878. Nth Magical Number

[zeikar]

Problem link

https://leetcode.com/problems/nth-magical-number/

Problem Summary

a 또는 b로 나누어 지는 수를 magical number라 할 때, n번째 magical number를 구하는 문제

Solution

일단 n이 크니까 탐색 범위를 줄여보자

LCM을 쓰면 a, b에 동시에 나누어지는 수를 구할 수 있고 이를 이용해 자르면 된다.
자른 다음 나머지 n에 대해 하나씩 a, b를 더해가면서 n번째 수를 구할 수 있다.

솔루션을 보면 이진 탐색으로 구하면 더 빠르게 구할 수 있다. x라는 수가 몇 번째인지 구할 수 있기 때문이다. (x // a + x // b - x // lcm). n번째의 최소의 x를 구하면 된다.

Source Code

class Solution:
    def lcm(self, a, b):
        return a // self.gcd(a, b) * b

    def gcd(self, a, b):
        if b == 0:
            return a
        return self.gcd(b, a % b)

    def nthMagicalNumber(self, n: int, a: int, b: int) -> int:
        lcm = self.lcm(a, b)
        cnt = lcm // a + lcm // b - 1

        skip = n // cnt
        remain = n % cnt
        ans = lcm * skip

        a_cnt = 1
        b_cnt = 1
        for _ in range(remain):
            if a * a_cnt < b * b_cnt:
                ans = lcm * skip + a * a_cnt
                a_cnt += 1
            else:
                ans = lcm * skip + b * b_cnt
                b_cnt += 1

        return ans % int(1e9 + 7)