A shifted geometric distribution with p=0.75 via bit twiddling

shifted_geom.py

# Came across the `randomLevel` function in
#
#   https://github.com/automerge/automerge/blob/main/backend/skip_list.js
#
# which implements a geometric distribution with bit twiddling.
#
# Below is a pedagogical implementation.

import random
from collections import Counter


def simplified():
    """
    Shifted geometric distribution with,
    
       Pr[X=x] = (1-p)^xp
    
    for,
    
      p = 0.75
    
    and,
    
      support = {1, 2, ..., 16}
    
    While technically truncated, 
     
      Pr[X > 16] = 2.328306436538698e-10
    
    so it's pretty much okay. 
    
    (You can get other distributions, too, if you are clever.)
    """
    
    
    # Generate a random uint32
    u = random.randint(0, 2**32-1)    
    
    # Divide the 32 bits into 16 bernoulli trials. 
    for x in range(1, 16+1):
        
        # Define the lower bound.
        v = 1 << (32 - 2 * x)
        
        # If u >= v at least one bit above v is set.
        #
        # Since we're doing this in 16 2-bit intervals we have,
        #
        #   Pr(HH) + Pr(HT) + Pr(TH) = 0.75
        #
        # for each loop trial, implementing the expansion.
    
        if u >= v:
            return x
        
    return 16
            
         
n = 1_000
counts = Counter(simplified() for i in range(n))
[counts.get(i, 0) / n for i in range(1, 16)]