markov-chain.py

import sys
from random import randint

def next_letter(probs, letter):
    if not letter in probs:
        return None
    probs_for_letter = probs[letter]
    sum_of_probs = 0
    for prob in probs_for_letter:
        sum_of_probs += probs_for_letter[prob]
    letter = randint(0, sum_of_probs-1)
    for prob in probs_for_letter:
        letter -= probs_for_letter[prob]
        if letter <= 0:
            return prob



filename = sys.argv[1]
generate_count = sys.argv[2]
length_from = sys.argv[3]
length_to = sys.argv[4]

# Read file with names, one name per line
names = []
with open(filename, 'r') as f:
    for line in f:
        names.append(line.strip())

# Process letter probabilities
probabilities = {}
first_letters = {}

for name in names:
    if name[0] in first_letters:
        first_letters[name[0]] = first_letters[name[0]] + 1
    else:
        first_letters[name[0]] = 1

    for idx in range(0, len(name)-1):
        first_letter = name[idx]
        second_letter = name[idx+1]
        if first_letter in probabilities:
            if second_letter in probabilities[first_letter]:
                probabilities[first_letter][second_letter] = probabilities[first_letter][second_letter] + 1
            else:
                probabilities[first_letter][second_letter] = 1
        else:
            probabilities[first_letter] = {}
            probabilities[first_letter][second_letter] = 1

# Generate names
for name in range(0, int(generate_count)):
    target_length = randint(int(length_from), int(length_to))
    # Choose first letter
    new_name = ''
    fl_sum = 0
    for fl in first_letters:
        fl_sum += first_letters[fl]
    fl_rnd = randint(0, fl_sum-1)
    for fl in first_letters:
        fl_rnd -= first_letters[fl]
        if fl_rnd <= 0:
            new_name = fl
            break
    i = 0
    while i < target_length:
        i += 1
        new_letter = next_letter(probabilities, new_name[i-1])
        if new_letter == '$':
            if len(new_name) >= int(length_from):
                break
            else:
                i -= 1
                continue
        new_name += new_letter
    print(new_name)