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
- code
class Solution:
def nthMagicalNumber(self, n: int, a: int, b: int) -> int:
L = a * b // gcd(a, b)
# how many mag nums < x
def magnums_belowx(x):
return x // a + x // b - x // L
l, r = 1, n * min(a, b)
while l < r:
mid = l + (r - l) // 2
if magnums_belowx(mid) < n:
l = mid + 1
else:
r = mid
return l % (10 ** 9 + 7)
- code TLE
class Solution:
def nthMagicalNumber(self, n: int, a: int, b: int) -> int:
if a > b: # let a be the small one
a, b = b, a
# 2, 3, 2+2, 2+2+2(3+3), 2+2+2+2, 3+3+3
# 1, 2, 3, 4, 5, 6
c = 10 ** 9 + 7
ans = lasta = a
lasta = lasta + a
lastb = b
if a == b:
lastb = lastb + b
for i in range(2, n+1):
if lasta < lastb:
ans = lasta
lasta = (lasta + a) % c
elif lasta > lastb:
ans = lastb
lastb = (lastb + b) % c
else:
ans = lasta
lasta = (lasta + a) % c
lastb = (lastb + b) % c
return ans