Solution to Decanting problem with capacities 8, 5, 3

decanting.py

"""
Problem Description:
You are given three vessels A, B and C of capacities 8, 5 and 3 gallons respectively. A is filled, while B and C are empty.
Divide the liquid in A into two equal quantities. 
(From Graph Theory, Narsingh Deo)
http://bit.ly/1Dnp4HA

Usage:
import decanting as dc
import networkx as nx
graph = dc.get_decanter_state_graph()
nx.shortest_path(graph, '8-0-0','4-4-0')

Output:
['8-0-0', '3-5-0', '3-2-3', '6-2-0', '6-0-2', '1-5-2', '1-4-3', '4-4-0']
"""

import itertools
import networkx as nx

transfer_permutations = list(itertools.permutations([0, 1, 2], 2))

def get_state_hash(state):
    return '-'.join([str(i) for i in state])

capacities = [8, 5, 3]

def get_next_decanter_states(current_state):
    next_states = []
    for start, end in transfer_permutations:
        end_remaining = capacities[end] - current_state[end]
        start_filled = current_state[start]
        if end_remaining == 0:
            continue
        if start_filled == 0:
            continue

        if start_filled >= end_remaining:
            fill_end_state = list(current_state)
            fill_end_state[start] -= end_remaining
            fill_end_state[end] += end_remaining
            
            next_states.append(fill_end_state)

        elif end_remaining > start_filled:
            empty_start_state = list(current_state)
            empty_start_state[start] -= start_filled
            empty_start_state[end] += start_filled
            
            next_states.append(empty_start_state)
            
    return next_states


def add_state_to_existing_states(state, existing_states):
    state_hash = get_state_hash(state)
    if not existing_states.get(state_hash):
        existing_states[state_hash] = state


def get_decanter_state_graph():
    current_state = [8, 0, 0]
    existing_states = {get_state_hash(current_state): current_state}

    state_queue = [current_state]

    graph = nx.DiGraph()

    iter_count = 0
    while True:

        iter_count += 1
        if iter_count > 1000:
            assert(False)

        if len(state_queue) == 0:
            break

        current_state = state_queue.pop(0)
        next_states = get_next_decanter_states(current_state)
        for state in next_states:
            if not existing_states.get(get_state_hash(state)):
                state_queue.append(state)
                add_state_to_existing_states(state, existing_states)
                graph.add_edge(get_state_hash(current_state), get_state_hash(state))

    return graph