gistfile1.txt

#!/usr/bin/env python3
import hashlib

"""
Write a function that buckets website users into experiment treatment groups with a weighted probability.

For example:

    Treatment A: 50%
    Treatment B: 20%
    Treatment C: 30%

When a user visits the site, they have a 30% chance of receiving treatment C.
"""

class User:
    def __init__(self, user_id=123, hash_value=None):
        self.user_id = user_id
        self.hash_value = hash_value

    def hash_int(self) -> int:
        return self.hash_value or int(hashlib.sha256(str(self.user_id).encode()).hexdigest(), 16)

def get_treatment_group(user: User):
    treatments = [
        ["A", 50],
        ["B", 20],
        ["C", 30],
    ]

    precision = sum(treatment[1] for treatment in treatments)

    value = user.hash_int() % precision

    threshold = 0
    for treatment, probability in treatments:
        threshold += probability

        if value <= threshold:
            return treatment

    raise Exception(f"Invalid data: {value}")

assert "A" == get_treatment_group(User())
assert "A" == get_treatment_group(User(hash_value=5))
assert "A" == get_treatment_group(User(hash_value=50))
assert "B" == get_treatment_group(User(hash_value=51))
assert "C" == get_treatment_group(User(hash_value=71))
assert "C" == get_treatment_group(User(hash_value=99))


assert "A" == get_treatment_group(User(hash_value=105))
assert "A" == get_treatment_group(User(hash_value=150))
assert "B" == get_treatment_group(User(hash_value=151))
assert "C" == get_treatment_group(User(hash_value=171))
assert "C" == get_treatment_group(User(hash_value=199))


"""
Follow up Qs:
    - how to make treatment the same if the user returns
    - how would you test this?
"""