AtCoder Beginner Contest 321 E

https://atcoder.jp/contests/abc321/tasks/abc321_e

2進で考えるとよいです。
始点から上と下に分けて考えます。上は一つ上がって、そこから今来たノードと違うノードの下を探ります。Kが大きければさらに上も見ます。下はKが大きすぎるときに注意が必要です。

// Complete Binary Tree
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_query() -> (i64, i64, i64) {
    let v = read_vec();
    let (N, X, K) = (v[0], v[1], v[2]);
    (N, X, K)
}

fn lower(N: i64, X: i64, K: i64) -> i64 {
    if K < 0 {
        return 0
    }
    
    for k in 0..K {
        if X << k > N {
            return 0
        }
    }
    let first = X << K;
    let last = (X + 1) << K;
    if N < first {
        0
    }
    else if N < last {
        N - first + 1
    }
    else {
        last - first
    }
}

fn upper(N: i64, X: i64, X_other: i64, K: i64) -> i64 {
    if X == 0 {
        0
    }
    else if K <= 0 {
        0
    }
    else if K == 1 {
        1
    }
    else {
        lower(N, (X << 2) + 1 - X_other, K - 2) + upper(N, X >> 1, X, K - 1)
    }
}

fn F_each(N: i64, X: i64, K: i64) -> i64 {
    lower(N, X, K) + upper(N, X >> 1, X, K)
}

fn F(T: usize) {
    for _ in 0..T {
        let (N, X, K) = read_query();
        println!("{}", F_each(N, X, K))
    }
}

fn main() {
    let T = read();
    F(T)
}