build_frequency_counts.py

import sys

from collections import defaultdict

def build_frequency_counts(text: str, n: int = 2) -> dict[tuple[list[str]], dict[str, int]]:
    if n < 1 or n > len(text):
        print(f"Error: n should be within 1 and {len(text)}, instead it is: {n}")
        sys.exit(1)

    counts = defaultdict(lambda: defaultdict(int)) # nested dict w/ a value set to 0

    for i in range(len(text) - n + 1):
        prefix = tuple(text[i : i + n - 1])
        next_char = text[i + n - 1]
        counts[prefix][next_char] += 1

    return counts

bigram_counts = build_frequency_counts("hellohello", n=2)
print("-" * 40)
print("text = \"hellohello\", n = 2")
print("-" * 40)
for prefix in bigram_counts:
    print(f"{prefix} → {dict(bigram_counts[prefix])}")
# ----------------------------------------
# text = "hellohello", n = 2
# ----------------------------------------
# ('h',) → {'e': 2}
# ('e',) → {'l': 2}
# ('l',) → {'l': 2, 'o': 2}
# ('o',) → {'h': 1}

trigram_counts = build_frequency_counts("hellohello", n=3)
print("-" * 40)
print("text = \"hellohello\", n = 3")
print("-" * 40)
for prefix in trigram_counts:
    print(f"{prefix} → {dict(trigram_counts[prefix])}")
# ----------------------------------------
# text = "hellohello", n = 3
# ----------------------------------------
# ('h', 'e') → {'l': 2}
# ('e', 'l') → {'l': 2}
# ('l', 'l') → {'o': 2}
# ('l', 'o') → {'h': 1}
# ('o', 'h') → {'e': 1}