Nth Magical Number

hard binary search math number theory

Problem

A positive integer is magical if it is divisible by a or by b. Given n, a, and b, return the nth smallest magical number, modulo 10⁹ + 7. The count of magical numbers ≤ x is x/a + x/b − x/lcm(a,b) (inclusion–exclusion), which is monotonic — so binary search for the smallest x whose count reaches n.

Inputn = 4, a = 2, b = 3

Output6

Magical numbers: 2, 3, 4, 6, … the 4th is 6.

from math import gcd

def nth_magical_number(n, a, b):
    lcm = a * b // gcd(a, b)
    lo, hi = 1, n * min(a, b)
    while lo < hi:
        mid = (lo + hi) // 2
        if mid // a + mid // b - mid // lcm >= n:
            hi = mid
        else:
            lo = mid + 1
    return lo % (10 ** 9 + 7)

function nthMagicalNumber(n, a, b) {
  const gcd = (x, y) => y ? gcd(y, x % y) : x;
  const lcm = a / gcd(a, b) * b;
  let lo = 1, hi = n * Math.min(a, b);
  while (lo < hi) {
    const mid = Math.floor((lo + hi) / 2);
    const cnt = Math.floor(mid / a) + Math.floor(mid / b) - Math.floor(mid / lcm);
    if (cnt >= n) hi = mid;
    else lo = mid + 1;
  }
  return lo % (1e9 + 7);
}

int nthMagicalNumber(int n, int a, int b) {
    long lcm = (long) a / gcd(a, b) * b;
    long lo = 1, hi = (long) n * Math.min(a, b);
    while (lo < hi) {
        long mid = (lo + hi) / 2;
        if (mid / a + mid / b - mid / lcm >= n) hi = mid;
        else lo = mid + 1;
    }
    return (int) (lo % 1_000_000_007);
}
private int gcd(int x, int y) { return y == 0 ? x : gcd(y, x % y); }

int nthMagicalNumber(int n, int a, int b) {
    long long lcm = (long long) a / __gcd(a, b) * b;
    long long lo = 1, hi = (long long) n * min(a, b);
    while (lo < hi) {
        long long mid = (lo + hi) / 2;
        if (mid / a + mid / b - mid / lcm >= n) hi = mid;
        else lo = mid + 1;
    }
    return lo % 1000000007;
}

Explanation

You could just list out magical numbers one by one until you reach the nth, but for large n that is far too slow. The clever move is to binary search on the answer itself instead of searching a list.

The key insight: for any number x, we can instantly count how many magical numbers are ≤ x. That count is x/a + x/b - x/lcm. We add the multiples of a and b, then subtract multiples of lcm(a, b) once because they got counted twice — this is inclusion-exclusion.

Because this count never decreases as x grows, we can binary search for the smallest x whose count reaches n. If count(mid) >= n the answer is at mid or lower, so we set hi = mid; otherwise we need a bigger value, so lo = mid + 1.

Example: n = 4, a = 2, b = 3, so lcm = 6. Test x = 6: 6/2 + 6/3 - 6/6 = 3 + 2 - 1 = 4, which reaches n = 4, and no smaller x does. So the 4th magical number is 6.

Each step halves the search window, so we only need about log checks, and the final answer is taken modulo 10^9 + 7.

Time: O(log(n·min(a,b))) Space: O(1)