Sample `count` indices without replacement. If `count >= n`, returns all indices (shuffled). Uses rejection sampling: resample if we pick an already-selected index. For count << n this is efficient. For count close to n, falls back to shuffle.
(&self, rng: &mut SeedableRng, count: usize)
| 98 | /// Uses rejection sampling: resample if we pick an already-selected index. |
| 99 | /// For count << n this is efficient. For count close to n, falls back to shuffle. |
| 100 | pub fn sample_without_replacement(&self, rng: &mut SeedableRng, count: usize) -> Vec<usize> { |
| 101 | let count = count.min(self.n); |
| 102 | |
| 103 | if count == self.n { |
| 104 | // Return all indices. |
| 105 | return (0..self.n).collect(); |
| 106 | } |
| 107 | |
| 108 | // For small count/n ratio, use rejection sampling. |
| 109 | if count <= self.n / 2 { |
| 110 | let mut selected = std::collections::HashSet::with_capacity(count); |
| 111 | let mut result = Vec::with_capacity(count); |
| 112 | let max_attempts = count * 20; // Prevent infinite loop on degenerate distributions. |
| 113 | let mut attempts = 0; |
| 114 | |
| 115 | while result.len() < count && attempts < max_attempts { |
| 116 | let idx = self.sample(rng); |
| 117 | if selected.insert(idx) { |
| 118 | result.push(idx); |
| 119 | } |
| 120 | attempts += 1; |
| 121 | } |
| 122 | result |
| 123 | } else { |
| 124 | // For large count, shuffle all indices and take first `count`. |
| 125 | let mut indices: Vec<usize> = (0..self.n).collect(); |
| 126 | // Fisher-Yates shuffle. |
| 127 | for i in (1..self.n).rev() { |
| 128 | let j = rng.gen_range((i + 1) as u64) as usize; |
| 129 | indices.swap(i, j); |
| 130 | } |
| 131 | indices.truncate(count); |
| 132 | indices |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | /// Sample `count` indices with replacement (same index can appear multiple times). |
| 137 | pub fn sample_with_replacement(&self, rng: &mut SeedableRng, count: usize) -> Vec<usize> { |