字节面试算法题输出化学式中每个原子的个数

2021.8.31 某同学的字节一面手撕算法题

给出类似于 HM2(H2ON3A)3N2 这样的类化学表达式
输出每个原子的个数

看起来有点意思

思路：

遇到左括号亚栈 有括号弹出栈
栈中存放固定大小为元素总数的数组，其记录每个元素的出现次数
同时主要到字母后面与右括号的数字，讲栈顶记录的出现次数与其相乘即可。

# 写上 所有 化学元素和 原子序数的对应 。这里比较懒先写几个用得到的
DIC = {
    'H': 0,
    'He': 1,
    'Li': 2,
    'C': 5,
    'O': 7,
    'S': 15,
    'Al': 12,
    'Cu': 29,
}

RE_DIC = {
    0: 'H',
    1: 'He',
    2: 'Li',
    5: 'C',
    7: 'O',
    15: 'S',
    12: 'Al',
    29: 'Cu',
}


def solve(s):
    stack = [[0] * 100]  # 化学元素先设定一百个吧。每个栈帧是个数组 数组按照元素顺序存放着元素的出现次数
    n = len(s)
    for idx, item in enumerate(s):
        if item.isupper(): # 化学元素名字总是一个大写字母开头后面都是小写字母

            # 支持多字符的元素，识别出类似于 Al Cu 这样的元素
            end = idx + 1
            while end <= n - 1 and s[end].islower():
                end += 1
            name = s[idx:end]

            # 记录右下缀 ,如 读出 Al2 中数字 2
            times = 1
            if end <= n - 1 and s[end].isdecimal():
                times = int(s[end])
            stack[-1][DIC[name]] += times

        elif item == '(': # 压栈
            stack.append([0] * 100)

        elif item == ')': # 出栈
            times = 1
            if idx != n - 1 and s[idx + 1].isdecimal():  # 如果 ) 后面是数字 要加倍
                times = int(s[idx + 1])
            popped = stack.pop()
            # 把弹出的信息，合并到栈顶中去
            for i in range(len(popped)):
                stack[-1][i] += popped[i] * times

    return stack


if __name__ == '__main__':
    # s = "Al2(SO4)3"
    s = "Cu2(OH)2CO3"
    ret = solve(s)
    for i in range(100):
        if ret[0][i] != 0:
            print(f"原子{RE_DIC[i]} : {ret[0][i]}")

输出

原子H : 2
原子C : 1
原子O : 5
原子Cu : 2