LeetCode: 878. Nth Magical Number

题目描述

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.

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

Example 4:

Input: N = 3, A = 6, B = 4
Output: 8

Note:

1 <= N <= 10^9
2 <= A <= 40000
2 <= B <= 40000

解题思路

枚举 A*B 之前所有可能的组合情况，然后根据 N 与它的关系，计算第 N 个数。

AC 代码

class Solution {

public:
    int nthMagicalNumber(int N, int A, int B) {

        set<long long int> AB;
        for(int i = 1; i <= B; ++i)
        {
            AB.insert(A*i);
        }
        for(int i = 1; i <= A; ++i)
        {
            AB.insert(B*i);
        }

        int t = N / AB.size();
        int c = N % AB.size();

        long long int ans = 0;
        while(t--)
        {
            ans += (*AB.rbegin());
            ans %= 1000000007;
        }
        cout << c << endl;

        auto iter = AB.begin();
        if(c){
            while(--c) ++iter;
            ans += (*iter);
            ans =ans % 1000000007;
        }

        return ans;
    }
};