878

878. Nth Magical Number (Hard)

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

Companies:
Amazon

Related Topics:
Math, Binary Search

Solution 1. Binary Search + Inclusion-exclusion Principle

L <= R template.

// OJ: https://leetcode.com/problems/nth-magical-number/
// Author: github.com/lzl124631x
// Time: O(log(N * min(A, B)))
// Space: O(1)
class Solution {
public:
    int nthMagicalNumber(int n, int a, int b) {
        long L = 2, R = (long)n * min(a, b), mod = 1e9 + 7;
        auto valid = [&](long M) { // returns `true` is the number of divisible numbers <= M is greater than or equal to n
            return M / a + M / b - M / (a * b / gcd(a, b)) >= n;
        };
        while (L <= R) {
            long M = L + (R - L) / 2;
            if (valid(M)) R = M - 1;
            else L = M + 1;
        }
        return L % mod;
    }
};

Or L < R template

// OJ: https://leetcode.com/problems/nth-magical-number/
// Author: github.com/lzl124631x
// Time: O(log(N * min(A, B)))
// Space: O(1)
class Solution {
public:
    int nthMagicalNumber(int n, int a, int b) {
        long L = 2, R = (long)n * min(a, b), mod = 1e9 + 7;
        auto valid = [&](long M) { // returns `true` is the number of divisible numbers <= M is greater than or equal to n
            return M / a + M / b - M / (a * b / gcd(a, b)) >= n;
        };
        while (L < R) {
            long M = L + (R - L) / 2;
            if (valid(M)) R = M;
            else L = M + 1;
        }
        return L % mod;
    }
};

Solution 2. Math

The pattern of magical numbers repeats.

Let L be the least common multiple of a and b. If x < L is magical, then x + L is magical. Because, say x = k * a, L = a * b / gcd(a, b), then x + L = [k + b / gcd(a,b)] * a. Example: a = 3, b = 11, L = 33, x = 6, x + L = 39.

There are M = L/a + L/b - 1 magical numbers <= L:

• L/a of them are divisible by a
• L/b of them are divisible by b
• 1 of them are divisible by both a and b.

So, for the first L numbers [1, L], there are M of them are magical. Adding L to each of them, we know that for the second L numbers [L+1, 2L], there are also M magical numbers. And so on.

Suppose n = M * q + r and r < M. The first L * q numbers contain M * q magical numbers.

For the remainder r magical numbers, we enumerate the smallest r numbers among a, 2a, 3a, ... and b, 2b, 3b, ....

Time complexity: The calculation of the rth magical number after q * M is O(M) = O(A + B).

// OJ: https://leetcode.com/problems/nth-magical-number/
// Author: github.com/lzl124631x
// Time: O(A + B)
// Space: O(1)
class Solution {
public:
    int nthMagicalNumber(int n, int a, int b) {
        long mod = 1e9 + 7, L = a * b / gcd(a, b), M = L / a + L / b - 1, q = n / M, r = n % M, ans = (long)q * L % mod;
        int val[2] = {0, 0};
        for (int i = 0; i < r; ++i) {
            auto [x, y] = val;
            if (x <= y) val[0] += a;
            if (y <= x) val[1] += b;
        }
        ans += min(val[0], val[1]);
        return ans % mod;
    }
};