Ugly Number III

medium binary search inclusion-exclusion math

Problem

An ugly number is a positive integer divisible by a, b, or c. Given four integers n, a, b, and c, return the nth ugly number.

Inputn = 4, a = 2, b = 3, c = 5

Output5

The sequence of ugly numbers is 2, 3, 4, 5, 6, 8, ... The 4th one is 5.

from math import gcd

def nth_ugly_number(n, a, b, c):
    ab = a * b // gcd(a, b)
    ac = a * c // gcd(a, c)
    bc = b * c // gcd(b, c)
    abc = ab * c // gcd(ab, c)

    def count(x):
        return (x // a + x // b + x // c
                - x // ab - x // ac - x // bc
                + x // abc)

    lo, hi = 1, 2 * 10**9
    while lo < hi:
        mid = (lo + hi) // 2
        if count(mid) < n:
            lo = mid + 1
        else:
            hi = mid
    return lo

function nthUglyNumber(n, a, b, c) {
  const gcd = (x, y) => y === 0 ? x : gcd(y, x % y);
  const lcm = (x, y) => x / gcd(x, y) * y;
  const ab = lcm(a, b), ac = lcm(a, c), bc = lcm(b, c);
  const abc = lcm(ab, c);
  const count = (x) =>
    Math.floor(x / a) + Math.floor(x / b) + Math.floor(x / c)
    - Math.floor(x / ab) - Math.floor(x / ac) - Math.floor(x / bc)
    + Math.floor(x / abc);

  let lo = 1, hi = 2e9;
  while (lo < hi) {
    const mid = Math.floor((lo + hi) / 2);
    if (count(mid) < n) lo = mid + 1;
    else hi = mid;
  }
  return lo;
}

class Solution {
    public int nthUglyNumber(int n, int a, int b, int c) {
        long ab = lcm(a, b), ac = lcm(a, c), bc = lcm(b, c);
        long abc = lcm(ab, c);
        long lo = 1, hi = 2_000_000_000L;
        while (lo < hi) {
            long mid = (lo + hi) / 2;
            long cnt = mid/a + mid/b + mid/c
                     - mid/ab - mid/ac - mid/bc + mid/abc;
            if (cnt < n) lo = mid + 1;
            else hi = mid;
        }
        return (int) lo;
    }
    private long lcm(long x, long y) { return x / gcd(x, y) * y; }
    private long gcd(long x, long y) { return y == 0 ? x : gcd(y, x % y); }
}

class Solution {
public:
    int nthUglyNumber(int n, int a, int b, int c) {
        long ab = lcm((long)a, b), ac = lcm((long)a, c), bc = lcm((long)b, c);
        long abc = lcm(ab, (long)c);
        long lo = 1, hi = 2000000000L;
        while (lo < hi) {
            long mid = (lo + hi) / 2;
            long cnt = mid/a + mid/b + mid/c
                     - mid/ab - mid/ac - mid/bc + mid/abc;
            if (cnt < n) lo = mid + 1;
            else hi = mid;
        }
        return (int) lo;
    }
private:
    long lcm(long x, long y) { return x / __gcd(x, y) * y; }
};

Explanation

Instead of generating ugly numbers one by one, we binary search on the answer itself. We ask: "what is the smallest value x such that there are at least n ugly numbers up to x?" That value is exactly the nth ugly number.

To do that we need a fast way to count ugly numbers <= x. We use inclusion-exclusion. x // a counts multiples of a, and similarly for b and c. But numbers divisible by two of them get counted twice, so we subtract the multiples of each pair's LCM (ab, ac, bc). Numbers divisible by all three got removed too often, so we add back x // abc.

This count only goes up as x grows, which is what makes binary search valid. If count(mid) < n we need a bigger value, so move lo up; otherwise mid might be the answer, so pull hi down to mid.

Example: n=4, a=2, b=3, c=5. The ugly sequence is 2, 3, 4, 5, .... The search converges on x = 5, because count(5) = 4 reaches the target while smaller values do not.

Since the search range can be huge but we halve it each step, this only takes about log of the range many counting calls.

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