Module c

ABC 400 C - 2^a b^2

refs: https://atcoder.jp/contests/abc400/tasks/abc400_c

各 $a$ について, $b$ の取りうる最大値を求める形式で解答できる.

$\left\lfloor\sqrt{\frac{X}{2^a}}\right\rfloor$ を正確に求められればかなり高速となるが算出誤差が生じうるため, 除算を含まない二分探索を用いるのが最も簡易に済む.

$a = 1$ のとき, 探索される範囲は $X = 2 \times 1^2$, $X = 2 \times 2^2$, $X = 2 \times 3^2$, … となる.

$a = 2$ のとき, 探索される範囲は $X = 2^2 \times 1^2$, $X = 2^2 \times 2^2$, $X = 2^2 \times 3^2$, … となる.

$a = 3$ 以上の奇数のとき, $X = 2^3 \times 1^2 = 2 \times (2 \times 1)^2$, $X = 2^3 \times 2^2 = 2 \times (2 \times 2)^2$, … と, $a = 1$ の探索範囲に含まれる.

$a = 4$ 以上の偶数のとき, $X = 2^4 \times 1^2 = 2^2 \times (2 \times 1)^2$, $X = 2^4 \times 2^2 = 2^2 \times (2 \times 2)^2$, … と, $a = 2$ の探索範囲に含まれる.

use proconio::input;

fn main() {
    input! {
        n: u64,
    }

    println!("{}", possible_b(n, 1) + possible_b(n, 2));
}

fn possible_b(n: u64, a: u32) -> u64 {
    bin_search(0, n, |b| calc_x(a, b) <= n)
}

fn bin_search(l: u64, r: u64, predicate: impl Fn(u64) -> bool) -> u64 {
    if r - l <= 1 {
        return l;
    }
    let mid = (l + r) / 2;
    if predicate(mid) {
        bin_search(mid, r, predicate)
    } else {
        bin_search(l, mid, predicate)
    }
}

fn calc_x(a: u32, b: u64) -> u64 {
    2u64.saturating_pow(a).saturating_mul(b.saturating_pow(2))
}