Problem statement

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

Solution 1

Imagine A = 6 and B = 10, then we have several first numbers: [6, 10, 12, 18, 20, 24, 30]. Note, than after they will be repeated with same pattern, added +30, +60 and so on. In general, there will be lcm(A, B)//A + lcm(A, B)//B - 1 numbers in group, where lcm is least common multiplier. Then we just find number with index N: find where in period it is and how many periods we make.

Complexity

Time complexity is O(A + B).

Code

class Solution:
    def nthMagicalNumber(self, N, A, B):
        lcm, Q = A*B//gcd(A,B), 10**9 + 7
        cand = sorted([A*i for i in range(1,lcm//A)] + [B*i for i in range(1,lcm//B+1)])
        m = len(cand)
        return (cand[(N-1)%m] + lcm*((N-1)//m)) % Q

Solution 2

There is also Binary search approach: answer for our problem is increasing, when we increase N: so we can evaluate in O(1) how many magical numbers we have before given number x, using inclusion-exclusion formula.

Complexity

Total complexity is O(log(N) + log(A + B)).

Code

class Solution:
    def nthMagicalNumber(self, N, A, B):
        lcm, Q = A*B//gcd(A,B), 10**9 + 7
        
        beg, end = 0, N * min(A,B)
        while beg < end:
            mid = (beg + end)//2
            if mid//A + mid//B - mid//lcm < N:
                beg = mid + 1
            else:
                end = mid
        
        return beg % Q