leetcode 726

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        element_info = []  # list of [element_name, number_of_element, number_of_parentheses]
        parentheses_stack = []
        i = 0
        while i < len(formula):
            if formula[i] == '(':
                parentheses_stack.append('(')
                i += 1
            elif formula[i] == ')':
                number_of_parentheses = len(parentheses_stack)
                # if formula[i + 1].isdigit():
                if i + 1 < len(formula) and formula[i + 1].isdigit():
                    res = ''
                    num = 0
                    for k in formula[i + 1:]:
                        if k.isdigit():
                            res += k
                            num += 1
                        else:
                            break
                    number_of_element = int(res)
                    for j, element in enumerate(element_info):
                        if element[2] == number_of_parentheses:
                            element_info[j][1] *= number_of_element
                            element_info[j][2] -= 1
                    parentheses_stack.pop()
                    i += num + 1
                else:
                    parentheses_stack.pop()
                    i += 1
            else:
                if i < len(formula) - 1:
                    res = find_yuanzi(formula, i, i + 1)
                    element_name, number_of_element, i = res
                    element_info.append([element_name, number_of_element, len(parentheses_stack)])
                else:
                    element_info.append([formula[i], 1, len(parentheses_stack)])
                    i += 1
        dic = {}
        for i in element_info:
            if i[0] not in dic:
                dic[i[0]] = i[1]
            else:
                dic[i[0]] += i[1]
        res = ''
        for i in sorted(dic):
            print(i)
            print(dic[i])
            if int(dic[i]) > 1:
                res += i + str(dic[i])
            else:
                res += i
        return res


def find_yuanzi(formula, i, j):
    if formula[j].isdigit():
        if j + 1 < len(formula) and formula[j + 1].isdigit():
            res = ''
            num = 0
            for k in formula[j:]:
                if k.isdigit():
                    res += k
                    num += 1
                else:
                    break
            return formula[i:j], int(res), j + num
        else:
            return formula[i:j], int(formula[j]), j + 1
    elif formula[j].islower():
        return find_yuanzi(formula, i, j + 1)
    else:
        return formula[i:j], 1, j


formula = "(H)"
so = Solution()
so.countOfAtoms(formula)