AtCoder Beginner Contest 269 C

https://atcoder.jp/contests/abc269/tasks/abc269_c

ABCのC問題を解いていきます。

Pythonだとgeneratorを使ってこんな感じで書けますが、

def gen_submask(N, L):
    if L == 1:
        yield 0
        if N == 1:
            yield 1
    else:
        L1 = L / 2
        L2 = L - L1
        ns1 = list(gen_submask(N & ((1 << L1) - 1), L1))
        ns2 = list(gen_submask(N >> L1, L2))
        for n2, n1 in product(ns2, ns1):
            yield n1 + (n2 << L1)

RustでIteratorを返すのは難しいので、Vecを返すことにします。

// Submask
#![allow(non_snake_case)]

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 num_digits(n: i64, b: i64) -> i64 {
    if n == 0 { 0 } else { num_digits(n / b, b) + 1 }
}

fn create_submasks(N: i64, L: i64) -> Vec<i64> {
    if L <= 1 {
        (0..(N+1)).collect()
    }
    else {
        let L1 = L / 2;
        let L2 = L - L1;
        let ns1 = create_submasks(N & ((1 << L1) - 1), L1);
        let ns2 = create_submasks(N >> L1, L2);
        let mut ns = Vec::<i64>::new();
        for n2 in ns2.iter() {
            for n1 in ns1.iter() {
                ns.push(n1 + (n2 << L1))
            }
        }
        return ns
    }
}

fn main() {
    let N: i64 = read();
    let L = num_digits(N, 2);
    for n in create_submasks(N, L).iter() {
        println!("{}", n)
    }
}