Kortit II

warning: unused variable: `a_plays_1`
  --> input/code.rs:66:60
   |
66 | ...k = Callback { callback: |a_plays_1: &Vec<u32>, b_plays_1: Option<&Vec<u32>>, a_plays_2: Option<&Vec<u32>>, b_plays_2: Option<&Vec<u32...
   |                              ^^^^^^^^^ help: if this is intentional, prefix it with an underscore: `_a_plays_1`
   |
   = note: `#[warn(unused_variables)]` on by default

warning: unused variable: `b_plays_1`
  --> input/code.rs:66:82
   |
66 | ...k: |a_plays_1: &Vec<u32>, b_plays_1: Option<&Vec<u32>>, a_plays_2: Option<&Vec<u32>>, b_plays_2: Option<&Vec<u32>>, total_cards: u32, ...
   |                              ^^^^^^^^^ help: if this is intentional, prefix it with an underscore: `_b_plays_1`

warning: unused variable: `a_plays_2`
  --> input/code.rs:66:112
   |
66 | ...ays_1: Option<&Vec<u32>>, a_plays_2: Option<&Vec<u32>>, b_plays_2: Option<&Vec<u32>>, total_cards: u32, draws: u32, a_points: u32, b_p...
   |                              ^^^^^^^^^ help: if this is...

use std::cell::RefCell;
use std::collections::HashSet;
use std::io;

const MOD: u128 = 10u128.pow(9) + 7;

fn main() {
    let stdin = io::stdin();

    let t: u32 = {
        let mut input: String = String::new();
        stdin.read_line(&mut input).unwrap();
        input.trim().parse().unwrap()
    };

    for _ in 0..t {
        let values: Vec<u32> = {
            let mut input: String = String::new();
            stdin.read_line(&mut input).unwrap();
            input.trim().split_whitespace().map(|x| x.parse::<u32>().unwrap()).collect()
        };

        let answer = generate_answer(values[0], values[1], values[2]);
        println!("{}", answer);
    }
}

fn generate_answer(total_cards: u32, a_points: u32, b_points: u32) -> u128 {
    let total_points = a_points + b_points;

    if total_points > total_cards || total_points > 0 && (a_points >= total_points || b_points >= total_points) {
        return 0;
    }

    let draws = total_cards - a_points - b_points;
    let mut combinations: RefCell<u32> = RefCell::new(0);

    // Cards [(total_cards - draws + 1), total_cards] are consumed for draws.
    // no-op

    // A wins a_points times.
    // When A uses card (total_cards - draws), it must win. a_points - 1 are still left to be chosen.
    let mut a_plays = Vec::<u32>::from([total_cards - draws]); // This vec has its elements in descending order (largest first).
    // Permutations are not yet taken into account, so drawing cards from largest to smallest is fine.

    let mut a_callback = Callback { callback: |a_plays: &Vec<u32>, _: Option<&Vec<u32>>, _: Option<&Vec<u32>>, _: Option<&Vec<u32>>, total_cards: u32, draws: u32, a_points: u32, b_points: u32, counter: &mut RefCell<u32>| {
        let mut b_plays = Vec::<u32>::new(); // This vec has its elements in the same order as a_plays.
        let mut b_used_cards = HashSet::new();

        let mut b_callback = Callback { callback: |a_plays_1: &Vec<u32>, b_plays_1: Option<&Vec<u32>>, _: Option<&Vec<u32>>, _: Option<&Vec<u32>>, total_cards: u32, draws: u32, a_points: u32, b_points: u32, counter: &mut RefCell<u32>| {
            let a_plays_2: Vec<u32> = {
                let mut a_plays_2_set: HashSet<u32> = (1..=(total_cards - draws)).collect();
                for a in a_plays_1 {
                    a_plays_2_set.remove(a);
                }
                a_plays_2_set.iter().map(|x| *x).collect()
            };

            let mut b_plays_2 = Vec::<u32>::new();
            let mut b_used_cards = HashSet::new(); // TODO: reuse the old set
            for b in b_plays_1.unwrap() {
                b_used_cards.insert(*b);
            }

            //println!("B lose callback: \n\tA plays 1: {:?}\n\tB plays 1: {:?}\n\tA plays 2: {:?}", a_plays_1, b_plays_1, a_plays_2);
            let mut b_win_callback = Callback { callback: |a_plays_1: &Vec<u32>, b_plays_1: Option<&Vec<u32>>, a_plays_2: Option<&Vec<u32>>, b_plays_2: Option<&Vec<u32>>, total_cards: u32, draws: u32, a_points: u32, b_points: u32, counter: &mut RefCell<u32>| {
                counter.replace_with(|x| -> u32 { *x + 1});
            }, total_cards, draws, a_points, b_points, counter};
            choose_b_win(&a_plays_1, &a_plays_2, &b_plays_1.unwrap(), &mut b_plays_2, &mut b_used_cards, total_cards, draws, a_points, b_points, &mut b_win_callback);
        }, total_cards, draws, a_points, b_points, counter};
        choose_b_lose(&a_plays, &mut b_plays, &mut b_used_cards, total_cards, draws, a_points, b_points, &mut b_callback);
    }, total_cards, draws, a_points, b_points, counter: &mut combinations};
    choose_a_win(a_points, &mut a_plays, &mut a_callback);

    let mut combinations: u128 = combinations.take() as u128;

    // The cards resulting in draws could be any card.
    combinations = (combinations % MOD) * (nCr(total_cards as u128, draws as u128) % MOD);

    let combinations_with_permutations = combinations * (factorial(total_cards as u128) % MOD);
    combinations_with_permutations % MOD
}

fn nCr(n: u128, k: u128) -> u128 {
    if k > n {
        0
    } else {
        factorial(n) / (factorial(k) * factorial(n - k))
    }
}

fn factorial(n: u128) -> u128 {
    (1..=n).product::<u128>()
}

struct Callback<'a> {
    callback: fn(&Vec<u32>, Option<&Vec<u32>>, Option<&Vec<u32>>, Option<&Vec<u32>>, u32, u32, u32, u32, &mut RefCell<u32>),
    total_cards: u32,
    draws: u32,
    a_points: u32,
    b_points: u32,
    counter: &'a mut RefCell<u32>
}

impl Callback<'_> {
    fn invoke(&mut self, plays_a_1: &Vec<u32>, plays_b_1: Option<&Vec<u32>>, plays_a_2: Option<&Vec<u32>>, plays_b_2: Option<&Vec<u32>>) {
        (self.callback)(plays_a_1, plays_b_1, plays_a_2, plays_b_2, self.total_cards, self.draws, self.a_points, self.b_points, self.counter);
    }
}

fn choose_a_win(a_points: u32, a_plays: &mut Vec<u32>, callback: &mut Callback) {
    if a_plays.len() >= a_points as usize {
        callback.invoke(a_plays, None, None, None);
        return;
    }

    let max = a_plays.last().unwrap() - 1;
    let min = 1 + a_points - a_plays.len() as u32;

    // The card must be between [min, max].
    for a_card in min..=max {
        a_plays.push(a_card);
        choose_a_win(a_points, a_plays, callback);
        a_plays.pop();
    }
}

fn choose_b_lose(a_plays: &Vec<u32>, b_plays: &mut Vec<u32>, b_used_cards: &mut HashSet<u32>, total_cards: u32, draws: u32, a_points: u32, b_points: u32, callback: &mut Callback) {
    if b_plays.len() >= a_plays.len() {
        callback.invoke(a_plays, Some(b_plays), None, None);
        return;
    }

    let max = a_plays[b_plays.len()] - 1;
    let mut range = 1..=max;

    if b_plays.len() == a_plays.len() - 1 && !b_used_cards.contains(&1) {
        range = 1..=1; // If this is the last card and 1 hasn't been used yet, it must be used now.
    }

    for b_card in range {
        if b_used_cards.contains(&b_card) {
            continue;
        }
        b_used_cards.insert(b_card);
        b_plays.push(b_card);
        choose_b_lose(a_plays, b_plays, b_used_cards, total_cards, draws, a_points, b_points, callback);
        b_plays.pop();
        b_used_cards.remove(&b_card);
    }
}


fn choose_b_win(a_plays_1: &Vec<u32>, a_plays_2: &Vec<u32>, b_plays_1: &Vec<u32>, b_plays_2: &mut Vec<u32>, b_used_cards: &mut HashSet<u32>, total_cards: u32, draws: u32, a_points: u32, b_points: u32, callback: &mut Callback) {
    if b_plays_2.len() >= b_points as usize {
        callback.invoke(a_plays_1, Some(b_plays_2), Some(a_plays_1), Some(b_plays_2));
        return;
    }

    let max = total_cards - draws;
    let min = a_plays_2[b_plays_2.len()] + 1;
    let range = min..=max;

    for b_card in range {
        if b_used_cards.contains(&b_card) {
            continue;
        }
        b_used_cards.insert(b_card);
        b_plays_2.push(b_card);
        choose_b_win(a_plays_1 , a_plays_2, b_plays_1, b_plays_2, b_used_cards, total_cards, draws, a_points, b_points, callback);
        b_plays_2.pop();
        b_used_cards.remove(&b_card);
    }
}

Test 1

Group: 1, 2, 3, 4, 5

Verdict: ACCEPTED

Test 2

Group: 2, 3, 4, 5

Verdict: ACCEPTED

Test 3

Group: 3, 4, 5

Verdict: TIME LIMIT EXCEEDED

Test 4

Group: 4, 5

Verdict: TIME LIMIT EXCEEDED

Test 5

Group: 5

Verdict: TIME LIMIT EXCEEDED

Test 6

Group: 5

Verdict: TIME LIMIT EXCEEDED