| author |  Justin Whear <justin.whear@gmail.com> 2022-08-28 04:19:51-07:00 |
| committer |  GitHub <noreply@github.com> 2022-08-28 14:19:51+03:00 |
| log | 5bb8c03697fce798a966feb131f6d906863047ae |
| tree | 29f549a0bdcaf00cbcca65e1b4502555a6b91cc5 |
| parent | 0f27836c218ee84f5fb7d4c47dd93a920663c037 |
| signature |  Signed by PGP key 4AEE18F83AFDEB23 |
std.random: add weightedIndex function
`weightedIndex` picks from a selection of weighted indices.2 files changed, 62 insertions(+), 0 deletions(-)
lib/std/rand.zig+36
| ... | ... | @@ -337,6 +337,42 @@ pub const Random = struct { |
| 337 | 337 | mem.swap(T, &buf[i], &buf[j]); |
| 338 | 338 | } |
| 339 | 339 | } |
| 340 | ||
| 341 | /// Randomly selects an index into `proportions`, where the likelihood of each | |
| 342 | /// index is weighted by that proportion. | |
| 343 | /// | |
| 344 | /// This is useful for selecting an item from a slice where weights are not equal. | |
| 345 | /// `T` must be a numeric type capable of holding the sum of `proportions`. | |
| 346 | pub fn weightedIndex(r: std.rand.Random, comptime T: type, proportions: []T) usize { | |
| 347 | // This implementation works by summing the proportions and picking a random | |
| 348 | // point in [0, sum). We then loop over the proportions, accumulating | |
| 349 | // until our accumulator is greater than the random point. | |
| 350 | ||
| 351 | var sum: T = 0; | |
| 352 | for (proportions) |v| { | |
| 353 | sum += v; | |
| 354 | } | |
| 355 | ||
| 356 | const point = if (comptime std.meta.trait.isSignedInt(T)) | |
| 357 | r.intRangeLessThan(T, 0, sum) | |
| 358 | else if (comptime std.meta.trait.isUnsignedInt(T)) | |
| 359 | r.uintLessThan(T, sum) | |
| 360 | else if (comptime std.meta.trait.isFloat(T)) | |
| 361 | // take care that imprecision doesn't lead to a value slightly greater than sum | |
| 362 | std.math.min(r.float(T) * sum, sum - std.math.epsilon(T)) | |
| 363 | else | |
| 364 | @compileError("weightedIndex does not support proportions of type " ++ @typeName(T)); | |
| 365 | ||
| 366 | std.debug.assert(point < sum); | |
| 367 | ||
| 368 | var accumulator: T = 0; | |
| 369 | for (proportions) |p, index| { | |
| 370 | accumulator += p; | |
| 371 | if (point < accumulator) return index; | |
| 372 | } | |
| 373 | ||
| 374 | unreachable; | |
| 375 | } | |
| 340 | 376 | }; |
| 341 | 377 | |
| 342 | 378 | /// Convert a random integer 0 <= random_int <= maxValue(T), |
lib/std/rand/test.zig+26
| ... | ... | @@ -445,3 +445,29 @@ test "CSPRNG" { |
| 445 | 445 | const c = random.int(u64); |
| 446 | 446 | try expect(a ^ b ^ c != 0); |
| 447 | 447 | } |
| 448 | ||
| 449 | test "Random weightedIndex" { | |
| 450 | // Make sure weightedIndex works for various integers and floats | |
| 451 | inline for (.{ u64, i4, f32, f64 }) |T| { | |
| 452 | var prng = DefaultPrng.init(0); | |
| 453 | const random = prng.random(); | |
| 454 | ||
| 455 | var proportions = [_]T{ 2, 1, 1, 2 }; | |
| 456 | var counts = [_]f64{ 0, 0, 0, 0 }; | |
| 457 | ||
| 458 | const n_trials: u64 = 10_000; | |
| 459 | var i: usize = 0; | |
| 460 | while (i < n_trials) : (i += 1) { | |
| 461 | const pick = random.weightedIndex(T, &proportions); | |
| 462 | counts[pick] += 1; | |
| 463 | } | |
| 464 | ||
| 465 | // We expect the first and last counts to be roughly 2x the second and third | |
| 466 | const approxEqRel = std.math.approxEqRel; | |
| 467 | // Define "roughly" to be within 10% | |
| 468 | const tolerance = 0.1; | |
| 469 | try std.testing.expect(approxEqRel(f64, counts[0], counts[1] * 2, tolerance)); | |
| 470 | try std.testing.expect(approxEqRel(f64, counts[1], counts[2], tolerance)); | |
| 471 | try std.testing.expect(approxEqRel(f64, counts[2] * 2, counts[3], tolerance)); | |
| 472 | } | |
| 473 | } |