Add nextArrayElementWeightedFunc method to Module

edjCase · Feb 14, 2024 · 164e7f3 · 164e7f3
1 parent eff3ef7
commit 164e7f3
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/PseudoRandomX.mo b/src/PseudoRandomX.mo
@@ -22,6 +22,7 @@ module Module {
         nextBufferElement : <T>(buffer : Buffer.Buffer<T>) -> T;
         nextArrayElement : <T>(array : [T]) -> T;
         nextArrayElementWeighted : <T>(array : [(T, Float)]) -> T;
+        nextArrayElementWeightedFunc : <T>(array : [T], weightFunc : (T) -> Float) -> T;
         shuffleBuffer : <T>(buffer : Buffer.Buffer<T>) -> ();
     };
 
@@ -153,6 +154,14 @@ module Module {
             Prelude.unreachable();
         };
 
+        public func nextArrayElementWeightedFunc<T>(array : [T], weightFunc : (T) -> Float) : T {
+            let weightedArray = Array.map<T, (T, Float)>(
+                array,
+                func(item : T) : (T, Float) = (item, weightFunc(item)),
+            );
+            nextArrayElementWeighted(weightedArray);
+        };
+
         public func shuffleBuffer<T>(buffer : Buffer.Buffer<T>) {
             let bufferSize = buffer.size();
             if (bufferSize == 0) {