exercise-12-3-6-pt-2.py

import re


class Equation:
    """
    A Class for representing a single equation in a system of equations.

    In the form:

    `2x_0 + 3x_1 - 4x_2 = 6`

    where `x_n` represents the nth index from the input variables.
    """

    def __init__(self, lhs: str, rhs: int):
        self.lhs: str = lhs
        self.rhs: int = rhs

    def evaluate(self, values: list[int]) -> int:
        assert len(self.lhs), "Cannot have empty expression for left-hand side!"

        # plug in values for `x_n`'s
        pattern = r"x_\d+"

        def get_next_replacement(match: re.Match[str]) -> str:
            group = match.group()
            return str(values[int(group[2:])])

        replaced_lhs = re.sub(
            pattern=pattern, repl=get_next_replacement, string=self.lhs
        )

        try:
            return eval(replaced_lhs)
        except:
            print(
                f"Could not evaluate {self.lhs=} with expected {self.rhs=} and {replaced_lhs=}"
            )
            exit(1)


class SystemOfEquations:
    """
    Represents a system of equations containing multiple `evaluate`-able
    equations.

    Check if values satisfies system in polynomial time.
    """

    def __init__(self):
        self.equations: list[Equation] = []

    def add(self, eq: Equation) -> None:
        self.equations.append(eq)

    def satisfies(self, values: list[int]) -> bool:
        """
        O(n * k)

        where n is the number of equations in the system and
        k is the number of values
        """
        return all([eq.evaluate(values) == eq.rhs for eq in self.equations])


if __name__ == "__main__":
    soe = SystemOfEquations()
    soe.add(Equation("x_0 + x_1 + x_2", 2))
    soe.add(Equation("x_0 + x_2", 1))
    soe.add(Equation("x_1 + x_2", 2))
    assert soe.satisfies([0, 1, 1])
    assert not soe.satisfies([2, 3, 4])
    print("✅")