CCF 201912-3 化学方程式 【Python 100分】 【LeetCode 726（hard）升级版】

现在CCF第三题就直接上LC hard了?

试题编号：	201912-3
试题名称：	化学方程式
时间限制：	1.0s
内存限制：	512.0MB

思路：栈+正则表达式

# coding=utf-8
# Copyright (c) 2020. ItsukiFujii.

import re
import collections

list_formula = []
list_terms_left = []
list_terms_right = []
list_res = []
dict_atoms_left = collections.defaultdict(int)
dict_atoms_right = collections.defaultdict(int)


# calculate each single compound
def fk_it(formula):
    stack = list()
    stack.append(collections.Counter())
    found = re.findall(r'([A-Z][a-z]?)(\d*)|(\()|(\))(\d*)', formula)
    for name, n_atom, open_paren, close_paren, coef in found:
        if name:
            stack[-1][name] += int(n_atom or 1)
        if open_paren:
            stack.append(collections.Counter())
        if close_paren:
            top = stack.pop()
            for key in top.keys():
                top[key] *= int(coef or 1)
            stack[-1] += top
    return dict(stack.pop())


# handle each single equation
def handle_equation(terms_all, dict_atoms):
    for term in terms_all:
        coef_term = re.match(r'\d+', term)
        coef_term = 1 if coef_term is None else int(coef_term.group())
        tmp = fk_it(term)
        for key in tmp.keys():
            tmp[key] *= coef_term
        for k, v in tmp.items():
            if str(dict_atoms.keys()).find(k) != -1:
                dict_atoms[k] += v
            else:
                dict_atoms.update({k: v})


if __name__ == '__main__':
    num_of_formula = int(input())
    for i in range(num_of_formula):
        list_formula.append(input())
        list_res.append('Y')

    # split the equations
    for item in list_formula:
        list_terms_left.append(item[:item.find('=')].split('+'))
        list_terms_right.append(item[item.find('=') + 1:].split('+'))

    # judge whether the number of atom on both side is equal
    cntr_tmp = 0
    for terms_all_left, terms_all_right in zip(list_terms_left, list_terms_right):
        dict_atoms_left.clear()
        dict_atoms_right.clear()

        # left part
        handle_equation(terms_all_left, dict_atoms_left)
        # right part
        handle_equation(terms_all_right, dict_atoms_right)
        # compare the numbers
        if dict_atoms_left != dict_atoms_right:
            list_res[cntr_tmp] = 'N'

        # print(dict_atoms_left, dict_atoms_right, list_res[cntr_tmp])  # for debug
        print(list_res[cntr_tmp])
        cntr_tmp += 1

用了regex真的很慢。。。取巧了一下，为了考试的时候节省时间。