Termial random using multiplication optimization

termial_bench_mult.py

#!/usr/bin/env python3

import math
import os
import random
import sys
import timeit

_ROUNDS = int(os.getenv("TERMIAL_DEFAULT_ROUNDS", 3))
_ITERATIONS = int(os.getenv("TERMIAL_ITERATIONS", 1000000))
_N = int(os.getenv("TERMIAL_N", 10000))


def random_termial_v0(n: int) -> int:
    """Get index to an item from a list of size n.

    The higher index of an item is, the lower probability of picking it.
    """
    termial_of_n = ((n * n) + n) >> 1
    choice = random.randrange(termial_of_n)
    i = math.floor(-0.5 + math.sqrt(0.25 + (choice << 1)))
    return n - 1 - i


def random_termial_v1(n: int) -> int:
    """Get index to an item from a list of size n.

    The higher index of an item is, the lower probability of picking it.
    """
    termial_of_n = ((n ** 2) + n) >> 1
    choice = random.randrange(termial_of_n)
    i = math.floor(-0.5 + math.sqrt(0.25 + (choice << 1)))
    return n - 1 - i


def random_termial_v2(n: int) -> int:
    """Get index to an item from a list of size n.

    The higher index of an item is, the lower probability of picking it.
    """
    if n % 2:  # or n & 1
        # odd numbers
        nearest = n + 1
        termial_of_n = ((1 + nearest) * (nearest / 2)) - nearest
    else:
        termial_of_n = (1 + n) * (n / 2)

    choice = random.randrange(termial_of_n)
    i = math.floor(-0.5 + math.sqrt(0.25 + (choice << 1)))
    return n - 1 - i


def random_termial_v3(n: int) -> int:
    """Get index to an item from a list of size n.

    The higher index of an item is, the lower probability of picking it.
    """
    addition = (n * (n & 1))
    nearest = n + addition
    termial_of_n = ((1 + nearest) * (nearest >> 1)) - addition
    choice = random.randrange(termial_of_n)
    i = math.floor(-0.5 + math.sqrt(0.25 + (choice << 1)))
    return n - 1 - i


def main():
    print("n          = ", _N)
    print("rounds     = ", _ROUNDS)
    print("iterations = ", _ITERATIONS)
    print()
    for i in range(_ROUNDS):
        print(f">>> Round {i}")
        for func in (random_termial_v0, random_termial_v1, random_termial_v2, random_termial_v3):
            print(
                "\t",
                func.__name__,
                ": ",
                timeit.timeit(
                    f"{func.__name__}({_N})",
                    number=_ITERATIONS,
                    globals=globals(),
                ),
            )


__name__ == "__main__" and main()

n          =  10000
rounds     =  3
iterations =  1000000

>>> Round 0
	 random_termial_v0 :  1.1368820359930396
	 random_termial_v1 :  1.2894328990369104
	 random_termial_v2 :  1.2646742290235125
	 random_termial_v3 :  1.2897816539625637
>>> Round 1
	 random_termial_v0 :  1.1827631939668208
	 random_termial_v1 :  1.3847550799837336
	 random_termial_v2 :  1.2822561039938591
	 random_termial_v3 :  1.460273706994485
>>> Round 2
	 random_termial_v0 :  1.369450077007059
	 random_termial_v1 :  1.6252998230047524
	 random_termial_v2 :  1.26978313800646
	 random_termial_v3 :  1.404374714998994

Based on response by @michalziman: https://medium.com/p/e39b9ca7aaa3/info