Stern-Brocot Treeで高速にN以下で表現される既約分数の近似値を求めたい

先日のABC333のG問題にて、Stern-Brocot Treeを用いる問題が出題されました。が、調べて実装してもTLEになってしまったのでその対策を行なった話をします。

Stern-Brocot Treeとは

自分も昨日はじめて耳にしました。
$[a/b,c/d]$の区間をノードにもっていて、子ノードに$[a/b,a+c/b+d]$,$[a+c/b+d,c/d]$を持つような木で、二分探索の要領でひとつずつ下に潜っていけばN以下で表現される既約分数の近似値がもとまるといったモノになっております。

詳しくはこちらの記事を参照されるとよいと思います。

TLEがとれない

愚直に実装したところ、だいたい0msや1msで実行できているのですが、なんか4ケースでTLEになってしまっていました。特定のケースの場合にめっちゃ時間がかかっている模様。

こういうときは境界付近を疑うとよさそうなので、めっちゃ小さい場合を考えてみます。
$[0/1,1/1]$ の子ノードは、
$[0/1,1/2]$ の子ノードは、
$[0/1,1/3]$ の子ノードは、
$[0/1,1/4]$ の子ノードは、
$[0/1,1/5]$ の子ノードは・・・

という感じで、0/1付近や1/1付近は一度の操作で区間に使用される数字が1しか増えないことに気づきました。一方、今回の制約は $N \leqq 10^{10}$ なので、0.000000000000001みたいな入力がくると間に合いません。

処理を改良

これまでは、以下のような感じで次のノードに行っていましたが、ここを頑張って高速化しました。

fn stern_brocot(
    pt: u128,
    qt: u128,
    pl: u128,
    ql: u128,
    pr: u128,
    qr: u128,
) -> (u128, u128, u128, u128) {
    let pm = pl + pr;
    let qm = ql + qr;
    if pt * qm < qt * pm {
        return (pl, ql, pm, qm);
    } else {
        return (pm, qm, pr, qr);
    }
}

$[a/b,c/d]$ -> $[a/b,a+c/b+d]$
ではなく、
$[a/b,c/d]$ -> $[a/b,ax+c/bx+d]$
となるような$x$をみつけて処理を高速化することを試みました。

子ノードに潜っていく方向を
右、左、左、左、右、右、右、右、右、左
だったものを
右×1、左×3、右×5、左×1
にするイメージです。

注意点として、右×1000000000みたいな場合に設定していた$N$をこえてしまうので適当なところで打ち切る処理をいれています。
右側に潜る場合と左側に潜る場合で計算がそれぞれ必要だったので、実装は少し面倒でした・・・が、これで高速化できて、全ケース1msで通るようになりました！！

fn stern_brocot(
    pt: u128,
    qt: u128,
    pl: u128,
    ql: u128,
    pr: u128,
    qr: u128,
    n: u128,
) -> (u128, u128, u128, u128) {
    let pm = pl + pr;
    let qm = ql + qr;
    if pt * qm < qt * pm {
        let x = ((pr * qt - pt * qr) / (pt * ql - pl * qt))
            .min((n - qr) / ql)
            .max(1);
        return (pl, ql, pl * x + pr, ql * x + qr);
    } else {
        let x = ((pt * ql - pl * qt) / (pr * qt - pt * qr))
            .min((n - ql) / qr)
            .max(1);
        return (pl + pr * x, ql + qr * x, pr, qr);
    }
}

最後までご覧いただきありがとうございました！
いいね押していただけるとモチベ上がりますので嬉しいです