Please read the new rule regarding the restriction on the use of AI tools. ×

Issue with k-Tree (431), probably some integer

Revision en1, by kadobot, 2021-01-03 12:43:51

I'm beginning with Rust, so probably my issue is related to that (maybe when summing the count or applying the modulo). Because I cannot spot where my logic can be wrong (test case #5 passes but #7 don't ?!).

https://codeforces.me/contest/431/submission/103079462

use std::{
    collections::HashMap,
    io::{self, BufRead},
};

fn main() {
    let stdin = io::stdin();
    let mut lines = stdin.lock().lines();
    let input = lines.next().unwrap().unwrap();
    let input = input
        .split_whitespace()
        .map(|n| n.parse::<u64>().unwrap())
        .collect::<Vec<_>>();
    let (n, k, d) = (input[0], input[1], input[2]);
    println!("{}", solve(n, k, d));
}

fn edges(id: u64, n: u64, k: u64, d: u64, drill: bool, dp: &mut HashMap<(u64, bool), u64>) -> u64 {
    if id == 0 {
        return 1;
    }

    if let Some(res) = dp.get(&(id, drill)) {
        return *res;
    }

    let mut count = 0;
    for j in 1..=k {
        if let Some(cur_id) = id.checked_sub(j) {
            if cur_id >= d {
                let new_bubble = if drill { drill } else { j >= d };
                let res = edges(cur_id, n, k, d, new_bubble, dp);
                dp.insert((cur_id, drill), res);
                count += res % (10e9 as u64 + 7);
            } else if j >= d || drill {
                count += edges(cur_id, n, k, d, true, dp);
            }
        }
    }

    dp.insert((id, drill), count);
    count
}

fn solve(n: u64, k: u64, d: u64) -> u64 {
    let mut dp = HashMap::new();
    edges(n, n, k, d, false, &mut dp)
}

mod tests {
    use crate::solve;

    #[test]
    fn example_1() {
        assert_eq!(3, solve(3, 3, 2));
    }

    #[test]
    fn example_2() {
        assert_eq!(1, solve(3, 3, 3));
    }

    #[test]
    fn example_3() {
        assert_eq!(6, solve(4, 3, 2));
    }

    #[test]
    fn example_4() {
        assert_eq!(7, solve(4, 5, 2));
    }

       #[test]
    fn test_case_5() {
        assert_eq!(110682188, solve(28, 6, 3));
    }

    #[test]
    fn test_case_7() {
        assert_eq!(295630102, solve(50, 6, 3));
    }
}

Tags rust, k-tree

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en3 English kadobot 2021-01-03 12:48:14 20
en2 English kadobot 2021-01-03 12:45:56 477 new version of code (not submitted)
en1 English kadobot 2021-01-03 12:43:51 2232 Initial revision (published)