AtCoder Beginner Contest 356 D

https://atcoder.jp/contests/abc356/tasks/abc356_d

ビットごとにカウントすると簡単です。

// Masked Popcount
#![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 ////////////////////

const D: u64 = 998244353;

fn read_input() -> (u64, u64) {
    let v: Vec<u64> = read_vec();
    let (N, M) = (v[0], v[1]);
    (N, M)
}

fn F(N: u64, M: u64) -> u64 {
    let mut counter: u64 = 0;
    for i in 0.. {
        if (N >> i) == 0 {
            break
        }
        if ((M >> i) & 1) == 0 {
            continue
        }
        
        let mut c: u64 = (N >> (i+1)) << i;
        if ((N >> i) & 1) == 1 {
            c += (N & ((1 << i) - 1)) + 1
        }
        counter = (counter + c) % D
    }
    return counter
}

fn main() {
    let (N, M) = read_input();
    println!("{}", F(N, M))
}