Pouring water problem

water_pour.py

"""
  Inital we have N empty cup with capacity(c1, c2, c3... cn)
  Posible move:
    empty one cup
    fill one cup
    pour from one cup to another
  We have unlimited water source, find a way to fill a target amount
"""

from collections import deque

class WaterPour:
    def __init__(self, capacity):
        self.capacity = capacity
        self.initial_state = (0,) * len(capacity)

    def update(self, state, index, value):
        return state[:index] + (value,) + state[index + 1:]

    def empty(self, state, glass):
        return self.update(state, glass, 0), f"Empty {glass}"

    def fill(self, state, glass):
        return self.update(state, glass, self.capacity[glass]), f"Fill {glass}"
    
    def pour(self, state, from_glass, to_glass):
        capacity = self.capacity
        amount = min(
            state[from_glass], capacity[to_glass] - state[to_glass])
        return self.update(
            self.update(state, from_glass, state[from_glass] - amount),
            to_glass,
            state[to_glass] + amount
        ), f"Pour from {from_glass}({state[from_glass]}) to {to_glass}({state[to_glass]})"

    def all_states(self, state):
        n = len(state)
        for g in range(n):
            yield self.empty(state, g)
            yield self.fill(state, g)
            for g2 in range(n):
                if g != g2:
                    yield(self.pour(state, g, g2))

    def trace(self, state, prev):
        paths = [(state, "Complete")]
        while state in prev:
            state, msg = prev[state]
            paths.append((state, msg))
        return paths[::-1]

    def solve(self, target):
        initial_state = self.initial_state
        queue = deque([initial_state])
        prev_states = {initial_state: (None, 'Initial')}

        while queue:
            state = queue.popleft()
            if target in state:
                return True, self.trace(state, prev_states)
            for new_state, msg in self.all_states(state):
                if new_state not in prev_states:
                    prev_states[new_state] = (state, msg)
                    queue.append(new_state)
        
        return False, None