Skip to content

Instantly share code, notes, and snippets.

@hatoo

hatoo/collatz_memo.rs

Created July 21, 2020 08:42
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save hatoo/22b3046f9d9304252dd059ff7621a668 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save hatoo/22b3046f9d9304252dd059ff7621a668 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. hatoo created this gist Jul 21, 2020.
    37 changes: 37 additions & 0 deletions collatz_memo.rs
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,37 @@
    #![allow(unused_macros)]
    #![allow(dead_code)]
    #![allow(unused_imports)]

    use std::io::stdin;
    use std::collections::{HashMap, BTreeMap};

    const U_INF: usize = 1 << 60;
    const I_INF: isize = 1 << 60;

    fn main() {
    let mut buf = String::new();
    stdin().read_line(&mut buf).unwrap();
    let n = buf.split_whitespace().next().unwrap().parse().unwrap();
    // let mut memo = HashMap::new();
    let mut memo = vec![None; n];


    let mut acc = 0;
    for i in 1..=n {
    acc += collatz(i, &mut memo);
    acc %= 1000000007;
    }
    println!("{} {}", acc, memo.len());
    }

    fn collatz(i: usize, memo: &mut [Option<usize>]) -> usize {
    if let Some(Some(v)) = memo.get(i) {
    return *v;
    }
    let cnt = if i == 1 { 0 } else if i % 2 == 0 { 1 + collatz(i / 2, memo) } else { 1 + collatz(i * 3 + 1, memo) };

    if let Some(m) = memo.get_mut(i) {
    *m = Some(cnt);
    }
    cnt
    }