nth-magical-number

Try it on leetcode

Description

A positive integer is magical if it is divisible by either a or b.

Given the three integers n, a, and b, return the nth magical number. Since the answer may be very large, return it modulo 109 + 7.

 

Example 1:

Input: n = 1, a = 2, b = 3
Output: 2

Example 2:

Input: n = 4, a = 2, b = 3
Output: 6

Example 3:

Input: n = 5, a = 2, b = 4
Output: 10

Example 4:

Input: n = 3, a = 6, b = 4
Output: 8

 

Constraints:

• 1 <= n <= 109
• 2 <= a, b <= 4 * 104

Solution(Python)

class Solution(object):
    def nthMagicalNumber(self, N, A, B):
        from math import gcd

        MOD = 10**9 + 7

        L = A // gcd(A, B) * B
        M = L // A + L // B - 1
        q, r = divmod(N, M)

        if r == 0:
            return q * L % MOD

        heads = [A, B]
        for _ in range(int(r) - 1):
            if heads[0] <= heads[1]:
                heads[0] += A
            else:
                heads[1] += B

        return (q * L + min(heads)) % MOD