competitive/tools/
xorshift.rs

1use std::{
2    collections::hash_map::RandomState,
3    hash::{BuildHasher, Hasher},
4    iter::repeat_with,
5};
6
7#[derive(Clone, Debug)]
8pub struct Xorshift {
9    y: u64,
10}
11
12impl Xorshift {
13    pub fn new_with_seed(seed: u64) -> Self {
14        Xorshift { y: seed }
15    }
16
17    pub fn new() -> Self {
18        Xorshift::new_with_seed(RandomState::new().build_hasher().finish())
19    }
20
21    pub fn rand64(&mut self) -> u64 {
22        self.y ^= self.y << 5;
23        self.y ^= self.y >> 17;
24        self.y ^= self.y << 11;
25        self.y
26    }
27
28    pub fn rand(&mut self, k: u64) -> u64 {
29        self.rand64() % k
30    }
31
32    pub fn rands(&mut self, k: u64, n: usize) -> Vec<u64> {
33        repeat_with(|| self.rand(k)).take(n).collect()
34    }
35
36    pub fn randf(&mut self) -> f64 {
37        const UPPER_MASK: u64 = 0x3FF0_0000_0000_0000;
38        const LOWER_MASK: u64 = 0x000F_FFFF_FFFF_FFFF;
39        let x = self.rand64();
40        let tmp = UPPER_MASK | (x & LOWER_MASK);
41        let result: f64 = f64::from_bits(tmp);
42        f64::from_bits(f64::to_bits(result - 1.0) ^ (x >> 63))
43    }
44
45    pub fn gen_bool(&mut self, p: f64) -> bool {
46        self.randf() < p
47    }
48
49    pub fn shuffle<T>(&mut self, slice: &mut [T]) {
50        let mut n = slice.len();
51        while n > 1 {
52            let i = self.rand(n as _) as usize;
53            n -= 1;
54            slice.swap(i, n);
55        }
56    }
57}
58
59impl Default for Xorshift {
60    fn default() -> Self {
61        Xorshift::new_with_seed(0x2b99_2ddf_a232_49d6)
62    }
63}