leetcode：1201. 丑数 III【二分 + 数学 + 容斥原理】

分析

丑树序列不能重复，所以是三个倍数之并
但并不好求，容斥原理展开
利用最小公倍数展开公式成为一个单调不减的fx
fx即为小于等于x的丑树的个数
二分找到第一个使得fx = n的x
这里需要注意一下是小于的话l = mid + 1
大于等于r = mid
细节一下

ac code

class Solution:
    def nthUglyNumber(self, n: int, a: int, b: int, c: int) -> int:
        # 数学题 + 容斥原理
        # 小于等于x的丑数个数设成f(x)


        # 两个数的最小公倍数
        def lcm(a, b):
            return (a * b) // gcd(a, b)

        # 三个数的最小公倍数
        def lcm3(a, b, c):
            return lcm(lcm(a, b), c)

        # 容斥原理：|A or B or C| = |A| + |B| + |C| - |A and B| - |B and C| - |C and A| + |A and B and C|
        def f(x):
            f1 = x // a + x // b + x // c
            f2 = x // lcm(a, b) + x // lcm(b, c) + x // lcm(c, a)
            f3 = x // lcm3(a, b, c)
            return f1 - f2 + f3
        
        # 二分找到第一个使得f(x) = n的x 
        # 因为fx单调不减
        l, r = 1, 2 * 10 ** 9
        while l < r:
            mid = (l + r) // 2
            if f(mid) < n:
                l = mid + 1
            else:
                r = mid
        
        return l
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

总结

容斥原理简化问题 二分常规解决