yukicoder No.2066 Simple Math !

floor sum のいい感じの練習問題！

• 問題へのリンク

問題概要

正整数 が与えられる。

非負整数 を用いて という形で表せる正の整数のうち、 番目に小さいものを求めよ。

( 個の入力ケースが与えられる)

制約

• 
• 

考えたこと

いかにも二分探索という問題。次の判定問題を考える。

---

以上 以下の整数のうち、 という形で表せる整数の個数を求めよ

(この値が 個以上である最小の が答え)

---

「 という形で表せる正\整数」をまともに数えるのは大変なので、できれば の組を数える方に帰着させたい。

そこで、 を最大公約数 で割っておくことにする。その状態で答えを求めて、最終的に答えに をかければよい。よって、 は互いに素であると仮定してよい。

を一意にするために

が互いに素であるならば、次のことが言える。

---

とするとき、整数 に対して、

を満たす整数 の値の組はただ つに定まる

---

よって、 の範囲で を動かしていき、 となるような の個数を足していけばよい。

と変形できる。 のとき、 個であるとわかる。これを合算していく。floor sum が使える。

このままだと の係数が負になっているので、次のように設定する。

• 
• 
• 
• 

こうして、floor sum を適用すればよい。計算量は対数の 2 乗オーダーとなる。

コード

#include <bits/stdc++.h>
using namespace std;

// sum_{i=0}^{n-1} floor((a * i + b) / m)
// O(log(n + m + a + b))
// __int128 can be used for T
template<class T> T floor_sum(T n, T a, T b, T m) {
    if (n == 0) return 0;
    T res = 0;
    if (a >= m) {
        res += n * (n - 1) * (a / m) / 2;
        a %= m;
    }
    if (b >= m) {
        res += n * (b / m);
        b %= m;
    }
    if (a == 0) return res;
    T ymax = (a * n + b) / m, xmax = ymax * m - b;
    if (ymax == 0) return res;
    res += (n - (xmax + a - 1) / a) * ymax;
    res += floor_sum(ymax, m, (a - xmax % a) % a, a);
    return res;
}

// calc #n that can be expressed n =  Px + Qy (P, Q is coprime)
// 0 <= n <= M
long long calc_num(__int128 P, __int128 Q, __int128 M) {
    __int128 mp = M / P;
    __int128 N = min(mp + 1, Q);
    __int128 a = P, b = M + Q - a * (N - 1);
    return floor_sum(N, a, b, Q) - 1;
}

int main() {
    int CASE;
    cin >> CASE;
    while (CASE--) {
        long long P, Q, K;
        cin >> P >> Q >> K;

        long long G = gcd(P, Q);
        P /= G, Q /= G;

        long long low = -1, high = 1LL<<50;
        while (high - low > 1) {
            long long M = (low + high) / 2;
            if (calc_num(P, Q, M) >= K) high = M;
            else low = M;
        }
        cout << high * G << endl;
    }
}