2019年12月CSP考试第三题化学方程式解法

2019年12月CSP考试第三题化学方程式解法

import sys
import re


def merge(m1, m2):
    for key in m1.keys():
        if key in m2.keys():
            m2[key] += m1[key]
        else:
            m2[key] = m1[key]
    return m2


def mul_formula(coef, formula):
    for key in formula.keys():
        formula[key] *= coef
    return formula


def equation(s):
    expr_left, expr_right = s.split("=")
    ret_l, ret_r = {}, {}
    for sub in expr_left.split("+"):
        coef, formula = expr(sub)
        formula = mul_formula(coef, formula)
        ret_l = merge(ret_l, formula)
    for sub in expr_right.split("+"):
        coef, formula = expr(sub)
        formula = mul_formula(coef, formula)
        ret_r = merge(ret_r, formula)
    if ret_l == ret_r:
        return "Y"
    else:
        return "N"


def expr(s):
    coef, s = get_coef(s)
    formula = get_formula(s)
    return coef, formula


def get_coef(s: str):
    if len(s) == 0:
        return 1, ""
    for i, c in enumerate(s):
        if c.isdigit() == False:
            if i == 0:
                return 1, s
            else:
                return int(s[:i]), s[i:]
    return int(s), ""


def get_formula(s: str):
    sp = s
    ret = {}
    while(len(sp)):
        term, sp = get_term(sp)
        coef, sp = get_coef(sp)
        term = mul_formula(coef, term)
        ret = merge(ret, term)
    return ret


def get_term(s: str):
    if s[0] == "(":
        i = 0
        lp = 0
        while True:
            if s[i] == "(":
                lp += 1
            if s[i] == ")":
                lp -= 1
            if lp == 0 and i != 0:
                break
            else:
                i+=1
        formula = s[1:i]
        remain = s[i+1:]
        return get_formula(formula), remain
    else:
        return get_element(s)


def get_element(s: str):
    if len(s) >= 2:
        if s[1].islower():
            return {s[:2]: 1}, s[2:]
    return {s[0]: 1}, s[1:]


if __name__ == "__main__":
    n = int(input())
    for i in range(n):
        eq = input()
        print(equation(eq))
    # print("H2+O2=H2O")
    # print(equation("H2+O2=H2O"))
    # print("2H2+O2=2H2O")
    # print(equation("2H2+O2=2H2O"))
    # print("H2+Cl2=2NaCl")
    # print(equation("H2+Cl2=2NaCl"))
    # print("H2+Cl2=2HCl")
    # print(equation("H2+Cl2=2HCl"))
    # print("CH4+2O2=CO2+2H2O")
    # print(equation("CH4+2O2=CO2+2H2O"))
    # print("CaCl2+2AgNO3=Ca(NO3)2+2AgCl")
    # print(equation("CaCl2+2AgNO3=Ca(NO3)2+2AgCl"))
    # print("3Ba(OH)2+2H3PO4=Ba3(PO4)2+6H2O")
    # print(equation("3Ba(OH)2+2H3PO4=Ba3(PO4)2+6H2O"))
    # print("3Ba(OH)2+2H3PO4=6H2O+Ba3(PO4)2")
    # print(equation("3Ba(OH)2+2H3PO4=6H2O+Ba3(PO4)2"))
    # print("4Zn+10HNO3=4Zn(NO3)2+NH4NO3+3H2O")
    # print(equation("4Zn+10HNO3=4Zn(NO3)2+NH4NO3+3H2O"))
    # print("4Au+8NaCN+2H2O+O2=4Na(Au(CN)2)+4NaOH")
    # print(equation("4Au+8NaCN+2H2O+O2=4Na(Au(CN)2)+4NaOH"))
    # print("Cu+As=Cs+Au")
    # print(equation("Cu+As=Cs+Au"))