diff --git a/src/modules/helpers/index.ts b/src/modules/helpers/index.ts index 3bdc2f68362..7c5ee9972cd 100644 --- a/src/modules/helpers/index.ts +++ b/src/modules/helpers/index.ts @@ -459,6 +459,55 @@ export class HelpersModule { return array[index]; } + /** + * Returns a weighted random element from the given array. Each element of the array should be an object with two keys `weight` and `value`. + * + * - Each `weight` key should be a number representing the probability of selecting the value, relative to the sum of the weights. Weights can be any positive float or integer. + * - Each `value` key should be the corresponding value. + * + * For example, if there are two values A and B, with weights 1 and 2 respectively, then the probability of picking A is 1/3 and the probability of picking B is 2/3. + * + * @template T The type of the entries to pick from. + * @param array Array to pick the value from. + * + * @example + * faker.helpers.weightedArrayElement([{ weight: 5, value: 'sunny' }, { weight: 4, value: 'rainy' }, { weight: 1, value: 'snowy' }]) // 'sunny', 50% of the time, 'rainy' 40% of the time, 'snowy' 10% of the time + * + * @since 8.0.0 + */ + weightedArrayElement( + array: ReadonlyArray<{ weight: number; value: T }> + ): T { + if (array.length === 0) { + throw new FakerError( + 'weightedArrayElement expects an array with at least one element' + ); + } + + if (!array.every((elt) => elt.weight > 0)) { + throw new FakerError( + 'weightedArrayElement expects an array of { weight, value } objects where weight is a positive number' + ); + } + + const total = array.reduce((acc, { weight }) => acc + weight, 0); + const random = this.faker.number.float({ + min: 0, + max: total, + precision: 1e-9, + }); + let current = 0; + for (const { weight, value } of array) { + current += weight; + if (random < current) { + return value; + } + } + + // In case of rounding errors, return the last element + return array[array.length - 1].value; + } + /** * Returns a subset with random elements of the given array in random order. * diff --git a/test/__snapshots__/helpers.spec.ts.snap b/test/__snapshots__/helpers.spec.ts.snap index 95ac4a03aa8..c5714ea91a3 100644 --- a/test/__snapshots__/helpers.spec.ts.snap +++ b/test/__snapshots__/helpers.spec.ts.snap @@ -177,6 +177,10 @@ exports[`helpers > 42 > uniqueArray > with array 1`] = ` ] `; +exports[`helpers > 42 > weightedArrayElement > with array 1`] = `"sunny"`; + +exports[`helpers > 42 > weightedArrayElement > with array with percentages 1`] = `"sunny"`; + exports[`helpers > 1211 > arrayElement > noArgs 1`] = `"c"`; exports[`helpers > 1211 > arrayElement > with array 1`] = `"!"`; @@ -368,6 +372,10 @@ exports[`helpers > 1211 > uniqueArray > with array 1`] = ` ] `; +exports[`helpers > 1211 > weightedArrayElement > with array 1`] = `"snowy"`; + +exports[`helpers > 1211 > weightedArrayElement > with array with percentages 1`] = `"snowy"`; + exports[`helpers > 1337 > arrayElement > noArgs 1`] = `"a"`; exports[`helpers > 1337 > arrayElement > with array 1`] = `"l"`; @@ -541,3 +549,7 @@ exports[`helpers > 1337 > uniqueArray > with array 1`] = ` "d", ] `; + +exports[`helpers > 1337 > weightedArrayElement > with array 1`] = `"sunny"`; + +exports[`helpers > 1337 > weightedArrayElement > with array with percentages 1`] = `"sunny"`; diff --git a/test/helpers.spec.ts b/test/helpers.spec.ts index 80c79e2f15b..6ef7c374eff 100644 --- a/test/helpers.spec.ts +++ b/test/helpers.spec.ts @@ -59,6 +59,20 @@ describe('helpers', () => { t.it('noArgs').it('with array', 'Hello World!'.split('')); }); + t.describe('weightedArrayElement', (t) => { + t.it('with array', [ + { weight: 5, value: 'sunny' }, + { weight: 4, value: 'rainy' }, + { weight: 1, value: 'snowy' }, + ]); + + t.it('with array with percentages', [ + { weight: 0.5, value: 'sunny' }, + { weight: 0.4, value: 'rainy' }, + { weight: 0.1, value: 'snowy' }, + ]); + }); + t.describe('arrayElements', (t) => { t.it('noArgs') .it('with array', 'Hello World!'.split('')) @@ -145,6 +159,94 @@ describe('helpers', () => { }); }); + describe('weightedArrayElement', () => { + it('should return a weighted random element in the array', () => { + const testArray = [ + { weight: 10, value: 'hello' }, + { weight: 5, value: 'to' }, + { weight: 3, value: 'you' }, + { weight: 2, value: 'my' }, + { weight: 1, value: 'friend' }, + ]; + const actual = faker.helpers.weightedArrayElement(testArray); + + expect(testArray.map((a) => a.value)).toContain(actual); + }); + + it('should return a weighted random element in the array using floats', () => { + const testArray = [ + { weight: 0.1, value: 'hello' }, + { weight: 0.05, value: 'to' }, + { weight: 0.03, value: 'you' }, + { weight: 0.02, value: 'my' }, + { weight: 0.01, value: 'friend' }, + ]; + const actual = faker.helpers.weightedArrayElement(testArray); + + expect(testArray.map((a) => a.value)).toContain(actual); + }); + + it('should return the only element in the array when there is only 1', () => { + const testArray = [{ weight: 10, value: 'hello' }]; + const actual = faker.helpers.weightedArrayElement(testArray); + + expect(actual).toBe('hello'); + }); + + it('should throw if the array is empty', () => { + expect(() => faker.helpers.weightedArrayElement([])).toThrowError( + new FakerError( + 'weightedArrayElement expects an array with at least one element' + ) + ); + }); + + it('should allow falsey values', () => { + const testArray = [{ weight: 1, value: false }]; + const actual = faker.helpers.weightedArrayElement(testArray); + expect(actual).toBe(false); + }); + + it('should throw if any weight is zero', () => { + const testArray = [ + { weight: 0, value: 'hello' }, + { weight: 5, value: 'to' }, + ]; + expect(() => + faker.helpers.weightedArrayElement(testArray) + ).toThrowError( + new FakerError( + 'weightedArrayElement expects an array of { weight, value } objects where weight is a positive number' + ) + ); + }); + + it('should throw if any weight is negative', () => { + const testArray = [ + { weight: -1, value: 'hello' }, + { weight: 5, value: 'to' }, + ]; + expect(() => + faker.helpers.weightedArrayElement(testArray) + ).toThrowError( + new FakerError( + 'weightedArrayElement expects an array of { weight, value } objects where weight is a positive number' + ) + ); + }); + + it('should not throw with a frozen array', () => { + const testArray = [ + { weight: 7, value: 'ice' }, + { weight: 3, value: 'snow' }, + ]; + const frozenArray = Object.freeze(testArray); + expect(() => + faker.helpers.weightedArrayElement(frozenArray) + ).to.not.throw(); + }); + }); + describe('arrayElements', () => { it('should return a subset with random elements in the array', () => { const testArray = ['hello', 'to', 'you', 'my', 'friend'];