Random Merkle Forest

import random
from sys import argv
from math import log2

forest = []

def insert_in_forest(item):
    if item % 1000 == 0:
        print("Processing item", item)

    # print("Structure of forest:", [len(x) for x in forest])
    prob_gen_new_tree = 1 /(1 + 1.2 * log2(1 + 0.5 * len(forest)))
    sample = random.choices([0, 1], weights=[1 - prob_gen_new_tree, prob_gen_new_tree], k=1)[0]

    if sample == 1:
        # print("Creating new tree")
        forest.append([item])
        return
    # print("Adding to old trees")
    min_len = min([len(x) for x in forest])
    weights = [0.9**(len(x) - min_len) for x in forest]
    sample = random.choices(
        list(range(len(forest))),
        weights=weights,
        k=1)[0]

    forest[sample].append(item)


if __name__ == "__main__":
    try:
        NUM_ITEMS = int(argv[1])
    except:
        print("Usage: python3 rmf.py <number of items>")
        exit(0)

    for i in range(NUM_ITEMS):
        insert_in_forest(i)

    print("Number of trees:", len(forest))
    print("Average size:", sum([len(x) for x in forest]) / len(forest))

    total_rmf = sum([len(x) - 1 for x in forest])
    total_mf = NUM_ITEMS - 1
    # rmf_sig = len(forest)
    # mf_sig = 1
    print("Total number of internal nodes in Merkle forest:", total_rmf)
    print("Total number of internal nodes in full merkle tree:", total_mf)
    print("Savings:", (total_mf - total_rmf) / (total_mf) * 100, "%")