in evals/registry/data/solve-for-variable/tools/solve.py [0:0]
def solve(self, location, mistake=None):
"""
Solve an equation for one of the variables in the RHS.
"location" is the location of the variable to solve for
(see Expression.var_location).
"mistake" is an array of False/True of the same size as "loc",
indicating at which step in the solving process will we be
adding a mistake.
"""
if mistake is None:
mistake = len(location) * [False]
else:
assert len(mistake) == len(location)
while location:
arg_pos = location[0]
correct = not mistake[0]
# In what follows, the RHS is never simplified
# (_simplify_minus or _simplify_in_situ),
# because the "location" depends on the shape of the RHS tree
# and would no longer be valid if the RHS changed shape.
# The RHS only loses one node at a time, at the same time
# that "location" loses the corresponding entry.
# for the unary minus,
# pass the minus to the other side of the equation
if self.right.op == "minus":
if correct:
self.left = Expression(Prio.UNARY, "minus", self.left)
_, self.left = self.left._simplify_minus()
self.right = self.right.args[0]
# for binary operators, the variable is in/under args[arg_pos],
# so pass args[1 - arg_pos] to the other side of the equation
elif self.right.op == "+":
self.left = Expression(
Prio.ADD_SUB, "-" if correct else "+", self.left, self.right.args[1 - arg_pos]
)
self.left._simplify_in_situ()
self.right = self.right.args[arg_pos]
elif self.right.op == "-":
if arg_pos == 0:
self.left = Expression(
Prio.ADD_SUB, "+" if correct else "-", self.left, self.right.args[1]
)
self.left._simplify_in_situ()
self.right = self.right.args[0]
else:
self.left = Expression(
Prio.ADD_SUB, "-" if correct else "+", self.right.args[0], self.left
)
self.left._simplify_in_situ()
self.right = self.right.args[1]
elif self.right.op == "*":
self.left = Expression(
Prio.MUL_DIV, "/" if correct else "*", self.left, self.right.args[1 - arg_pos]
)
self.right = self.right.args[arg_pos]
elif self.right.op == "/":
if arg_pos == 0:
self.left = Expression(
Prio.MUL_DIV, "*" if correct else "/", self.left, self.right.args[1]
)
self.right = self.right.args[0]
else:
self.left = Expression(
Prio.MUL_DIV, "/" if correct else "*", self.right.args[0], self.left
)
self.right = self.right.args[1]
# consume this element and continue with the next RHS node
location = location[1:]
mistake = mistake[1:]
# iteration end; everything else was passed to the other side
# leaving the RHS with the single variable we wanted,
# so just switch the sides
self.left, self.right = self.right, self.left