From 09e23e3ed6294cb51ea4345258f9b514191d1af5 Mon Sep 17 00:00:00 2001
From: Subhomoy Haldar <hungrybluedev@gmail.com>
Date: Fri, 14 Oct 2022 04:54:02 +0100
Subject: [PATCH] rand: add rand.element and prng.element functions with unit
 test (#16068)

---
 vlib/rand/rand.v                | 15 +++++++++++++++
 vlib/rand/random_numbers_test.v | 18 ++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/vlib/rand/rand.v b/vlib/rand/rand.v
index cafd76fab4..288b53b8ff 100644
--- a/vlib/rand/rand.v
+++ b/vlib/rand/rand.v
@@ -383,6 +383,15 @@ pub fn (mut rng PRNG) choose<T>(array []T, k int) ?[]T {
 	return results
 }
 
+// element returns a random element from the given array.
+// Note that all the positions in the array have an equal chance of being selected. This means that if the array has repeating elements, then the probability of selecting a particular element is not uniform.
+pub fn (mut rng PRNG) element<T>(array []T) ?T {
+	if array.len == 0 {
+		return error('Cannot choose an element from an empty array.')
+	}
+	return array[rng.intn(array.len)!]
+}
+
 // sample samples k elements from the array with replacement.
 // This means the elements can repeat and the size of the sample may exceed the size of the array.
 pub fn (mut rng PRNG) sample<T>(array []T, k int) []T {
@@ -607,6 +616,12 @@ pub fn choose<T>(array []T, k int) ?[]T {
 	return default_rng.choose<T>(array, k)
 }
 
+// element returns a random element from the given array.
+// Note that all the positions in the array have an equal chance of being selected. This means that if the array has repeating elements, then the probability of selecting a particular element is not uniform.
+pub fn element<T>(array []T) ?T {
+	return default_rng.element<T>(array)
+}
+
 // sample samples k elements from the array with replacement.
 // This means the elements can repeat and the size of the sample may exceed the size of the array.
 pub fn sample<T>(array []T, k int) []T {
diff --git a/vlib/rand/random_numbers_test.v b/vlib/rand/random_numbers_test.v
index 5bda97c000..1329c4521b 100644
--- a/vlib/rand/random_numbers_test.v
+++ b/vlib/rand/random_numbers_test.v
@@ -411,3 +411,21 @@ fn test_sample() {
 		assert element in a
 	}
 }
+
+fn test_element1() {
+	a := ['one', 'two', 'four', 'five', 'six', 'seven']
+	for _ in 0 .. 30 {
+		e := rand.element(a)?
+		assert e in a
+		assert 'three' != e
+	}
+}
+
+fn test_element2() {
+	for _ in 0 .. 30 {
+		e := rand.element([1, 2, 5, 6, 7, 8])?
+		assert e in [1, 2, 5, 6, 7, 8]
+		assert 3 != e
+		assert 4 != e
+	}
+}