Dragon merging

math3444274.py

#!/usr/bin/python3

# https://math.stackexchange.com/q/3444274

from fractions import Fraction


def build_states(target_dragon):
    states = {}
    def toBinary(x, width):
        return [x & 1] + toBinary(x >> 1, width - 1) if width > 0 else []

    # States where no dragon size has two individuals
    for i in range(2 ** (target_dragon - 1)):
        states[tuple(toBinary(i, target_dragon - 1) + [0])] = len(states)
    # States where one dragon size has two individuals. We never have two sizes.
    for doubled_dragon in range(target_dragon - 1):
        for j in range(2 ** (target_dragon - 2)):
            state = toBinary(j, target_dragon - 2)
            state.insert(doubled_dragon, 2)
            # Optimisation: we can't get substrings [1, 1, 2] because the two is reached by merging
            # one of the two previous ones, and that leaves them on 0.
            if doubled_dragon > 1 and state[doubled_dragon - 2:doubled_dragon] == [1, 1]:
                continue
            states[tuple(state + [0])] = len(states)
    # Terminal state
    states[tuple([0] * (target_dragon - 1) + [1])] = len(states)
    print("Enumerated states:", len(states))
    return states


def spawn_matrix(states, probSpawn1):
    n = len(states)
    target_dragon = len(next(iter(states)))
    probSpawn2 = 1 - probSpawn1

    matrix = [[0] * n for _ in range(n)]
    for state, i in states.items():
        # We only spawn if there are no doubles
        if i >= 2 ** (target_dragon - 1):
            matrix[i][i] = 1
            continue

        spawn1 = list(state)
        spawn1[0] += 1
        matrix[states[tuple(spawn1)]][i] = probSpawn1

        spawn2 = list(state)
        spawn2[1] += 1
        matrix[states[tuple(spawn2)]][i] = probSpawn2
    return matrix


def merge_matrix_sparse(states, merge_dragon, probMerge1):
    n = len(states)
    target_dragon = len(next(iter(states)))
    probMerge2 = 1 - probMerge1

    matrix = [{} for _ in range(n)]
    for state, i in states.items():
        # We only merge if we have a double
        if state[merge_dragon] < 2:
            matrix[i][i] = 1
            continue

        # Special-case if we're one below the target
        if merge_dragon == target_dragon - 2:
            matrix[len(states) - 1][i] = 1
            continue

        merge1 = list(state)
        merge1[merge_dragon] -= 2
        merge1[merge_dragon + 1] += 1
        matrix[states[tuple(merge1)]][i] = probMerge1

        merge2 = list(state)
        merge2[merge_dragon] -= 2
        merge2[merge_dragon + 2] += 1
        matrix[states[tuple(merge2)] if merge_dragon + 2 < target_dragon - 1 else len(states) - 1][i] = probMerge2

    return matrix


def matrix_mul_sparse_dense(a, b):
    rows = len(a)
    cols = len(b[0])
    c = [[0] * cols for _ in range(rows)]
    for i in range(rows):
        for j in range(cols):
            for k, aik in a[i].items():
                c[i][j] += aik * b[k][j]
    return c


def to_sparse(matrix):
    sparse = []
    for row in matrix:
        sparse_row = {}
        for idx, val in enumerate(row):
            if val:
                sparse_row[idx] = val
        sparse.append(sparse_row)
    return sparse


def optimise(states, matrix, impossible_double_dragon):
    # Eliminate states with two of a given dragon size.
    # This should speed up processing.
    deleted_indices = set(idx for state, idx in states.items() if state[impossible_double_dragon] == 2)
    new_matrix = []
    for row_idx, row_contents in enumerate(matrix):
        if row_idx in deleted_indices:
            continue

        new_row = []
        for col_idx, col_value in enumerate(row_contents):
            if col_idx in deleted_indices:
                continue
            new_row.append(col_value)
        new_matrix.append(new_row)

    # We assume that all deleted_indices are contiguous
    cutoff = max(deleted_indices)
    new_states = {}
    for state, idx in states.items():
        if idx in deleted_indices:
            continue
        if idx > cutoff:
            idx -= len(deleted_indices)
        new_states[state] = idx
    return (new_states, new_matrix)


if __name__ == "__main__":
    probSpawn1 = Fraction(8, 10)
    probMerge1 = Fraction(85, 100)
    target_dragon = 3

    states = build_states(target_dragon)
    transition_matrix = spawn_matrix(states, probSpawn1)
    print("Applied spawn")
    for merge_dragon in range(target_dragon - 1):
        merge = merge_matrix_sparse(states, merge_dragon, probMerge1)
        transition_matrix = matrix_mul_sparse_dense(merge, transition_matrix)

        (states, transition_matrix) = optimise(states, transition_matrix, merge_dragon)
        print("Applied merge for", merge_dragon + 1, "and reduced to", len(states), "states")

    transition_matrix = to_sparse(transition_matrix)
    distrib = [[1]] + [[0]] * (len(states) - 1)
    previous_termination_prob = 0
    partial_expected = 0
    for t in range(1, 600):
        distrib = matrix_mul_sparse_dense(transition_matrix, distrib)
        termination_prob = distrib[-1][0]
        partial_expected += t * (termination_prob - previous_termination_prob)
        print(t, termination_prob, partial_expected)
        previous_termination_prob = termination_prob