1use std::f64;
2use std::iter::Iterator;
3
4pub struct PRNG {
5 pub coins: Box<dyn Iterator<Item = bool>>,
6}
7
8impl PRNG {
9 pub fn new(coins: Box<dyn Iterator<Item = bool>>) -> PRNG {
10 PRNG { coins }
11 }
12
13 pub fn draw(&mut self) -> f64 {
14 let mut bits = Vec::new();
15 for _ in 0..32 {
16 bits.push(self.coins.next().unwrap());
17 }
18 assert_eq!(bits.len(), 32);
19
20 let mut out = 0;
21 for &b in &bits {
22 out = (out << 1) | (b as u32);
23 }
24
25 let res = out as f64 / ((1u64 << 32) - 1) as f64;
26 assert!(res >= 0.0 && res <= 1.0);
27 res
28 }
29}
30
31pub fn afc(i: i32) -> f64 {
32 if i < 0 {
33 panic!("i should not be < 0");
34 } else if i == 0 {
35 return 0.0;
36 }
37
38 let frac_12 = 1.0 / 12.0;
39 let frac_360 = 1.0 / 360.0;
40 let frac_pi = 0.5 * f64::consts::PI.ln();
41
42 let i = i as f64;
43 (i + 0.5) * i.ln() - i + frac_12 / i - frac_360 / (i * i * i) + frac_pi
44}
45
46fn loggam(x: f64) -> f64 {
47 let a = [
48 8.333333333333333e-02,
49 -2.777777777777778e-03,
50 7.936507936507937e-04,
51 -5.952380952380952e-04,
52 8.417508417508418e-04,
53 -1.917526917526918e-03,
54 6.410256410256410e-03,
55 -2.955065359477124e-02,
56 1.796443723688307e-01,
57 -1.39243221690590,
58 ];
59
60 let mut x0 = x;
61 let mut n = 0;
62
63 if x0 == 1.0 || x0 == 2.0 {
64 return 0.0;
65 } else if x0 <= 7.0 {
66 n = (7.0 - x0) as i32;
67 x0 += n as f64;
68 }
69
70 let x2 = 1.0 / (x0 * x0);
71 let xp = 2.0 * f64::consts::PI;
72
73 let mut gl0 = a[9];
74 for k in (0..=8).rev() {
75 gl0 = gl0 * x2 + a[k];
76 }
77
78 let mut gl = gl0 / x0 + 0.5 * xp.ln() + (x0 - 0.5) * x0.ln() - x0;
79 if x <= 7.0 {
80 for _k in 1..=n {
81 gl -= (x0 - 1.0).ln();
82 x0 -= 1.0;
83 }
84 }
85 gl
86}
87
88pub fn hypergeometric_hyp(prng: &mut PRNG, good: f64, bad: f64, sample: f64) -> i32 {
89 let d1 = bad + good - sample;
90 let d2 = f64::min(bad, good);
91
92 let mut y = d2;
93 let mut k = sample;
94 while y > 0.0 {
95 let u = prng.draw();
96 y -= (u + y / (d1 + k)).floor();
97 k -= 1.0;
98 if k == 0.0 {
99 break;
100 }
101 }
102
103 let mut z = (d2 - y) as i32;
104 if good > bad {
105 z = sample as i32 - z;
106 }
107 z
108}
109
110pub fn hypergeometric_hrua(prng: &mut PRNG, good: f64, bad: f64, sample: f64) -> i32 {
111 const D1: f64 = 1.7155277699214135;
112 const D2: f64 = 0.8989161620588988;
113
114 let mingoodbad = good.min(bad);
115 let popsize = good + bad;
116 let maxgoodbad = good.max(bad);
117 let m = sample.min(popsize - sample);
118 let d4 = mingoodbad / popsize;
119 let d5 = 1.0 - d4;
120 let d6 = m * d4 + 0.5;
121 let d7 = ((popsize - m) * sample * d4 * d5 / (popsize - 1.0) + 0.5).sqrt();
122 let d8 = D1 * d7 + D2;
123 let d9 = ((m + 1.0) * (mingoodbad + 1.0) / (popsize + 2.0)).floor() as i32;
124 let d10 = loggam(d9 as f64 + 1.0)
125 + loggam(mingoodbad - d9 as f64 + 1.0)
126 + loggam(m - d9 as f64 + 1.0)
127 + loggam(maxgoodbad - m + d9 as f64 + 1.0);
128 let d11 = m.min(mingoodbad) + 1.0_f64.min((d6 + 16.0 * d7).floor());
129
130 loop {
131 let x = prng.draw();
132 let y = prng.draw();
133 let w = d6 + d8 * (y - 0.5) / x;
134
135 if w < 0.0 || w >= d11 {
136 continue;
137 }
138
139 let z = w.floor() as i32;
140 let t = d10
141 - (loggam(z as f64 + 1.0)
142 + loggam(mingoodbad - z as f64 + 1.0)
143 + loggam(m - z as f64 + 1.0)
144 + loggam(maxgoodbad - m + z as f64 + 1.0));
145
146 if (x * (4.0 - x) - 3.0) <= t {
147 let mut result = z;
148 if good > bad {
149 result = m as i32 - result;
150 }
151 if m < sample {
152 result = good as i32 - result;
153 }
154 return result;
155 }
156
157 if x * (x - t) >= 1.0 {
158 continue;
159 }
160
161 if 2.0 * x.ln() <= t {
162 let mut result = z;
163 if good > bad {
164 result = m as i32 - result;
165 }
166 if m < sample {
167 result = good as i32 - result;
168 }
169 return result;
170 }
171 }
172}
173
174pub type Coins = Box<dyn Iterator<Item = bool> + 'static>;
175
176pub fn rhyper(kk: i32, nn1: f64, nn2: f64, coins: Coins) -> i32 {
177 let mut prng = PRNG::new(coins);
178
179 if kk > 10 {
180 hypergeometric_hrua(&mut prng, nn1, nn2, kk as f64)
181 } else {
182 hypergeometric_hyp(&mut prng, nn1, nn2, kk as f64)
183 }
184}