7.4 Tree Diagrams, Tables, and Outcomes

from itertools import product
from fractions import Fraction

# A multistage experiment: flip a coin, then roll a die. The tree has two
# levels; this is the same information as a table you can read straight through.
# EDIT these two: use nonempty lists of distinct choices within each stage.
# Choices may be letters, numbers or whole words; counts are derived from them.
stage1 = ["H", "T"]
stage2 = [1, 2, 3, 4, 5, 6]

# Repeated labels would count the same named outcome more than once.
for name, choices in (("stage1", stage1), ("stage2", stage2)):
    if not choices:
        raise ValueError(f"{name} must be nonempty: include at least one choice.")
    if len(set(choices)) != len(choices):
        raise ValueError(f"{name} must contain distinct choices: remove repeated labels.")

space = list(product(stage1, stage2))
print(f"sample space: {len(stage1)} x {len(stage2)} = {len(space)} outcomes")
print()
# Ordered pairs keep the two stages separate, even with multiword choices.
table = [["stage 1 / stage 2"] + [repr(d) for d in stage2]]
for face in stage1:
    table.append([repr(face)] + [repr((face, d)) for d in stage2])
widths = [max(len(row[i]) for row in table) for i in range(len(table[0]))]
for row in table:
    print(" | ".join(cell.ljust(width) for cell, width in zip(row, widths)))

print()
p = Fraction(1, len(space))
print("Assume the stages are independent and each listed choice is equally likely.")
print(f"Then every outcome has probability {p}.")

# An EVENT is a grouping of outcomes, so its probability is just a count.
# The events are DERIVED from whatever the two stages now hold, so they stay
# meaningful after an edit instead of silently describing the old experiment.
first, last = stage1[0], stage2[-1]
events = {
    f"first stage is {first}":        lambda a, b: a == first,
    f"second stage is {last}":        lambda a, b: b == last,
    f"first is {first} AND second is {last}": lambda a, b: a == first and b == last,
    f"first is {first} OR second is {last}":  lambda a, b: a == first or b == last,
}
print()
for name, rule in events.items():
    hits = [o for o in space if rule(*o)]
    print(f"  P({name:<34}) = {len(hits)}/{len(space)} = {Fraction(len(hits), len(space))}")

print()
print(f"Adding another first-stage draw multiplies the number of outcomes by {len(stage1)}.")
print("The product counts the branches without listing every outcome.")
from math import prod
for stages in (2, 3, 4):
    sizes = [len(stage1)] * stages + [len(stage2)]
    print(f"  {stages} draws from the first stage ({len(stage1)} choices) then one "
          f"from the second ({len(stage2)}): "
          f"{' x '.join(str(s) for s in sizes)} = {prod(sizes)} outcomes")