#819364 (Rust) No.2136 Dice Calendar?

• 問題: No.2136 Dice Calendar?
• ユーザー: CuriousFairy315
• 提出日時: 2022-11-25 21:15:01
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 808 ms / 5,000 ms
• コード長: 7,644 bytes
• コンパイル時間: 14,710 ms
• コンパイル使用メモリ: 399,796 KB
• 実行使用メモリ: 71,424 KB
• 最終ジャッジ日時: 2024-10-02 03:46:09
• 合計ジャッジ時間: 17,749 ms
• ジャッジサーバーID （参考情報）: judge4 / judge1

ファイルパターン	結果
sample	AC * 3
other	AC * 24

ソースコード

const TRUE: &bool = &true;
const FALSE: &bool = &false;

#[derive(Clone, Debug)]
/// Efficient bool collection
pub struct BitSet {
    buf: Vec<u64>,
    size: usize,
}

impl BitSet {
    #[allow(dead_code)]
    pub fn new(size: usize) -> BitSet {
        BitSet {
            buf: vec![0; (size + 63) / 64],
            size,
        }
    }

    #[allow(dead_code)]
    pub fn set(&mut self, i: usize, b: bool) {
        assert!(i < self.size);
        if b {
            self.buf[i >> 6] |= 1 << (i & 63);
        } else {
            self.buf[i >> 6] &= !(1 << (i & 63));
        }
    }

    #[allow(dead_code)]
    pub fn count_ones(&self) -> u32 {
        self.buf.iter().map(|x| x.count_ones()).sum()
    }

    #[allow(dead_code)]
    fn chomp(&mut self) {
        let r = self.size & 63;
        if r != 0 {
            if let Some(x) = self.buf.last_mut() {
                let d = 64 - r;
                *x = (*x << d) >> d;
            }
        }
    }
}

impl std::ops::Index<usize> for BitSet {
    type Output = bool;
    fn index(&self, index: usize) -> &bool {
        [FALSE, TRUE][(self.buf[index >> 6] >> (index & 63)) as usize & 1]
    }
}

#[allow(clippy::suspicious_op_assign_impl)]
impl std::ops::ShlAssign<usize> for BitSet {
    fn shl_assign(&mut self, x: usize) {
        let q = x >> 6;
        let r = x & 63;

        if q >= self.buf.len() {
            for x in &mut self.buf {
                *x = 0;
            }
            return;
        }

        if r == 0 {
            for i in (q..self.buf.len()).rev() {
                self.buf[i] = self.buf[i - q];
            }
        } else {
            for i in (q + 1..self.buf.len()).rev() {
                self.buf[i] = (self.buf[i - q] << r) | (self.buf[i - q - 1] >> (64 - r));
            }
            self.buf[q] = self.buf[0] << r;
        }

        for x in &mut self.buf[..q] {
            *x = 0;
        }

        self.chomp();
    }
}

impl std::ops::Shl<usize> for BitSet {
    type Output = Self;

    fn shl(mut self, x: usize) -> Self {
        self <<= x;
        self
    }
}

#[allow(clippy::suspicious_op_assign_impl)]
impl std::ops::ShrAssign<usize> for BitSet {
    fn shr_assign(&mut self, x: usize) {
        let q = x >> 6;
        let r = x & 63;

        if q >= self.buf.len() {
            for x in &mut self.buf {
                *x = 0;
            }
            return;
        }

        if r == 0 {
            for i in 0..self.buf.len() - q {
                self.buf[i] = self.buf[i + q];
            }
        } else {
            for i in 0..self.buf.len() - q - 1 {
                self.buf[i] = (self.buf[i + q] >> r) | (self.buf[i + q + 1] << (64 - r));
            }
            let len = self.buf.len();
            self.buf[len - q - 1] = self.buf[len - 1] >> r;
        }

        let len = self.buf.len();
        for x in &mut self.buf[len - q..] {
            *x = 0;
        }
    }
}

impl std::ops::Shr<usize> for BitSet {
    type Output = Self;

    fn shr(mut self, x: usize) -> Self {
        self >>= x;
        self
    }
}

impl<'a> std::ops::BitAndAssign<&'a BitSet> for BitSet {
    fn bitand_assign(&mut self, rhs: &'a Self) {
        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
            *a &= *b;
        }
    }
}

impl<'a> std::ops::BitAnd<&'a BitSet> for BitSet {
    type Output = Self;
    fn bitand(mut self, rhs: &'a Self) -> Self {
        self &= rhs;
        self
    }
}

impl<'a> std::ops::BitOrAssign<&'a BitSet> for BitSet {
    fn bitor_assign(&mut self, rhs: &'a Self) {
        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
            *a |= *b;
        }
        self.chomp();
    }
}

impl<'a> std::ops::BitOr<&'a BitSet> for BitSet {
    type Output = Self;
    fn bitor(mut self, rhs: &'a Self) -> Self {
        self |= rhs;
        self
    }
}

impl<'a> std::ops::BitXorAssign<&'a BitSet> for BitSet {
    fn bitxor_assign(&mut self, rhs: &'a Self) {
        for (a, b) in self.buf.iter_mut().zip(rhs.buf.iter()) {
            *a ^= *b;
        }
        self.chomp();
    }
}

impl<'a> std::ops::BitXor<&'a BitSet> for BitSet {
    type Output = Self;
    fn bitxor(mut self, rhs: &'a Self) -> Self {
        self ^= rhs;
        self
    }
}

//proconio
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let mut s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};

    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };

    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };

    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };

    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };

    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}

const MOD: u64 = 998_244_353;

#[allow(unused_mut)]
fn main () {
    input! {
        n: usize,
        s: [[usize; 6]; n],
    }
    let factorial: Vec<u64> = (0..n + 1).scan(1u64, |cum, x| {
    	if x != 0 {*cum *= x as u64}
    	Some(*cum)
    }).collect();
    
    let mut set = BitSet::new(1 << n + 9);
    let mut now: Vec<u32> = Vec::with_capacity(3200000);
	let mut next_collection: Vec<u32> = Vec::with_capacity(3200000);
    now.push(0b111111111);
	for dice in s.iter() {
		let dice: Vec<usize> = dice.iter().map(|i| i - 1).collect();
		for multiset in now.iter() {
			next(*multiset, &dice, &mut set, &mut next_collection);
		}
		std::mem::swap(&mut now, &mut next_collection);
		next_collection.clear();
	}
	let ans = now.iter().fold(0u64, |cum, x| (cum + multichoose(&factorial, *x)) % MOD);
	println!("{}", ans);
}

fn calc_partition(multiset: u32) -> u64 { // 与えられた多重集合に対して、立っているbitの位置を保持する数列Pを返す
	let mut multiset = multiset;
	let mut partition = 0u64;
	for i in (5..50).step_by(5) {
		let lob = multiset.trailing_zeros();
		partition += 1u64 + (lob as u64) << i;
		multiset -= 1 << lob;
	}
	partition
}

fn get_partition(partition: u64, index: usize) -> usize {
	 // multiSetでindex番目に立っているbitの位置を求める
	 (partition >> 5 * index & 0b11111) as usize
}

fn multichoose(factorial: &Vec<u64>, multiset: u32) -> u64{ // multiSetで与えられた多重集合を並べてできる組合せ
	let partition = calc_partition(multiset);
	let mut multichoose = factorial[get_partition(partition, 9) - 9];
	for i in 0..9 {multichoose /= factorial[get_partition(partition, i + 1) - get_partition(partition, i) - 1]};
	multichoose
}

fn next(multiset: u32, dice: &Vec<usize>, set: &mut BitSet, stack: &mut Vec<u32>) { // diceを追加したときにできる新たな多重集合のうち、新しく発見したものをstackに入れる
	let partition = calc_partition(multiset);
	for result in dice {
		let mask = (1u32 << get_partition(partition, *result)) - 1;
		let next = (multiset & !mask) << 1 | multiset & mask;
		if !set[next as usize] {
			set.set(next as usize, true);
			stack.push(next);
		}
	}
}