@@ -389,28 +389,44 @@ where
389
389
390
390
let mut candidates = Vec :: with_capacity ( amount. as_usize ( ) ) ;
391
391
let mut index = N :: zero ( ) ;
392
- while index < length {
392
+ while index < length && candidates . len ( ) < amount . as_usize ( ) {
393
393
let weight = weight ( index. as_usize ( ) ) . into ( ) ;
394
394
if weight > 0.0 {
395
395
let key = rng. random :: < f64 > ( ) . ln ( ) / weight;
396
- if candidates. len ( ) < amount. as_usize ( ) {
397
- candidates. push ( Element { index, key } ) ;
398
- } else if key > candidates[ 0 ] . key {
399
- candidates[ 0 ] = Element { index, key } ;
400
- }
401
- candidates. sort_unstable ( ) ;
396
+ // let key = rng.random::<f64>().powf(1.0 / weight);
397
+ candidates. push ( Element { index, key } ) ;
402
398
} else if !( weight >= 0.0 ) {
403
399
return Err ( WeightError :: InvalidWeight ) ;
404
400
}
405
401
406
402
index += N :: one ( ) ;
407
403
}
404
+ candidates. sort_unstable ( ) ;
408
405
409
- let avail = candidates. len ( ) ;
410
- if avail < amount. as_usize ( ) {
406
+ if candidates. len ( ) < amount. as_usize ( ) {
411
407
return Err ( WeightError :: InsufficientNonZero ) ;
412
408
}
413
409
410
+ let mut x = rng. random :: < f64 > ( ) . ln ( ) / candidates[ 0 ] . key ;
411
+ while index < length {
412
+ let weight = weight ( index. as_usize ( ) ) . into ( ) ;
413
+ if weight > 0.0 {
414
+ x -= weight;
415
+ if x <= 0.0 {
416
+ let t = core:: f64:: consts:: E . powf ( candidates[ 0 ] . key * weight) ;
417
+ let key = rng. random_range ( t..1.0 ) . ln ( ) / weight;
418
+ candidates[ 0 ] = Element { index, key } ;
419
+ candidates. sort_unstable ( ) ;
420
+
421
+ x = rng. random :: < f64 > ( ) . ln ( ) / candidates[ 0 ] . key ;
422
+ }
423
+ } else if !( weight >= 0.0 ) {
424
+ return Err ( WeightError :: InvalidWeight ) ;
425
+ }
426
+
427
+ index += N :: one ( ) ;
428
+ }
429
+
414
430
Ok ( IndexVec :: from ( candidates. iter ( ) . map ( |elt| elt. index ) . collect ( ) ) )
415
431
}
416
432
0 commit comments