Search code examples
c++algorithmbinary-search

Intuition behind this solution


I am trying to solve a question on LeetCode.com:

A positive integer is magical if it is divisible by either A or B.
Return the N-th magical number. Since the answer may be very large, return it modulo 10^9 + 7. For e.g., for N = 5, A = 2, B = 4, output is 10.

The highly upvoted solution suggests 'searching' for the answer. But I do not understand how this could be modeled as a search problem - that too a binary search problem. Could someone please point out the intuition behind the same?

int nthMagicalNumber(int N, int A, int B) {
    long lcm = A * B / __gcd(A, B), l = 2, r = 1e14, mod = 1e9 + 7;
    while (l < r) {
        long m = (l + r) / 2;
        if (m / A + m / B - m / lcm < N) l = m + 1;
        else r = m;
    }
    return l % mod;
}

Thanks!


Solution

  • For a given number m, we can find out how many magical numbers are smaller than m: add m/A (the count of numbers divisible by A that are smaller than m) and m/B (the count of numbers divisible by B that are smaller than m), subtract m/lcm (the count of numbers divisible by both A and B that are smaller than m, to avoid double-counting them). Here lcm is the least common multiple of A and B; all numbers divisible by both A and B are multiples of lcm.

    From here, it's easy to use binary search. Start with a large enough interval (proving that the initial range is large enough based on the input bounds provided in the problem statement is left as an exercise for the reader). Find the middle m, figure out the count of magical numbers smaller than that middle. If that count is larger than N, then m overshoots and we need further consider the left half. Otherwise, m undershoots and we search the right half.


    This approach can be improved. Given sequences X_i=i*A and Y_i=i*B, we consider sequence Z that is an ordered union of X and Y; our task is to find Z_N. It is easy to see that Z can be described by a repeating pattern: after L=lcm(A, B) is reached, the same pattern repeats itself shifted up by L, then again shifted up by 2*L and so on.

    For example, if A=2 and B=3, the sequence Z looks like this:

    2 3 4 6 | 8 9 10 12 | 14 15 16 18 | ...
    

    It's easy to see a pattern of four numbers, repeating itself over and over, each time shifted by 6 = L = LCM(2, 3)

    The length of this pattern in the general case is M = L/A + L/B - 1: the count of numbers from sequence X that are <= L, plus the count of numbers from sequence Y that are <= L, minus 1 because L itself is counted twice. Thus, Z_{i+k*M} = Z_i + k*L.

    With this in mind, to find Z_N it is sufficient to find Z_{N modulo M}. For that, it's sufficient to apply binary search to the range [1, L] - to search just the first occurrence of the repeating pattern, since the rest of the sequence can be easily derived from that.