在问 · 算法题解

726. 原子的数量

困难栈与堆栈 + 哈希表力扣原题 ↗
在线写代码并运行看动画推演问 AI 讲解

题目描述

输入字符串 formula，它是一个合法化学式。统计其中每种原子的总个数。

原子名由一个大写字母开头，后接零个或多个小写字母。原子名后面可以跟一个数字表示个数，个数为 1 时不写数字。几个化学式直接拼接仍是化学式；一段化学式用括号括起来后可以再跟一个数字，表示括号内的整体重复这么多次。

返回一个字符串：按原子名字典序依次写出每种原子，名字后紧跟它的总数；总数为 1 时只写名字。

示例 1
输入：formula = "H2O"
输出："H2O"

解释：H 有 2 个，O 有 1 个。

示例 2
输入：formula = "Mg(OH)2"
输出："H2MgO2"

解释：H 2 个、Mg 1 个、O 2 个，按字典序为 H2MgO2。

示例 3
输入：formula = "K4(ON(SO3)2)2"
输出："K4N2O14S4"

解释：展开括号后 K 4 个、N 2 个、O 14 个、S 4 个。

数据范围

  • 1 <= formula.length <= 1000
  • formula 只含英文字母、数字、( 和 )。
  • formula 一定是合法的化学式。

解题思路：栈 + 哈希表

用一个“计数器栈”处理嵌套：栈中每层是一个哈希表，记录这一层括号里的原子个数。遇到 ( 压入一个空表；遇到原子就读出名字和后面的数字，累加到栈顶表；遇到 ) 就读出它后面的倍数 mult，弹出栈顶表，把里面每个原子的个数乘以 mult 后合并到新的栈顶。扫完后栈底就是总计数，按名字排序拼接即可。

为什么正确

栈的第 d 层记录的是“当前尚未闭合的第 d 层括号里，已经扫过的部分”的原子个数；栈底是括号外的部分。

复杂度分析

时间 O(n·d + m log m)，d 为括号嵌套深度（每次合并可能搬动一层的全部原子），m 为原子种类数；空间 O(n)。

易错点

原子名可能是多个字母（Mg、Be），数字也可能是多位（Be32），都要循环读完；数字省略时个数为 1。输出时个数为 1 不写数字。括号后的倍数要乘到括号内的每一种原子上，而不是只乘最后一个。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        stack = [collections.Counter()]
        i, n = 0, len(formula)
        while i < n:
            if formula[i] == '(':
                stack.append(collections.Counter())
                i += 1
            elif formula[i] == ')':
                i += 1
                j = i
                while i < n and formula[i].isdigit():
                    i += 1
                mult = int(formula[j:i] or 1)
                top = stack.pop()
                for name, c in top.items():
                    stack[-1][name] += c * mult
            else:
                j = i
                i += 1
                while i < n and formula[i].islower():
                    i += 1
                name = formula[j:i]
                j = i
                while i < n and formula[i].isdigit():
                    i += 1
                stack[-1][name] += int(formula[j:i] or 1)
        cnt = stack[0]
        return ''.join(name + (str(cnt[name]) if cnt[name] > 1 else '') for name in sorted(cnt))

载入编辑器，自己跑一遍 →

其他解法

栈 + 哈希表和递归下降是同一种分层合并的两种实现，都要逐层搬动计数；从右向左扫描先知道倍数，可以一遍直接累加到结果。

解法 2：递归下降解析 时间 O(n·d + m log m)，空间 O(n)

写一个 parse 函数，从当前位置读到右括号或末尾为止，返回这一段的原子计数。遇到原子就读名字和数字并累加。遇到左括号就递归调用 parse 得到子计数，跳过右括号后读倍数，乘上后合并。

把显式栈换成了系统调用栈，结构和文法一一对应，更容易扩展到更复杂的文法。嵌套很深时受递归深度限制。

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        n = len(formula)
        i = 0

        def read_num():
            nonlocal i
            j = i
            while i < n and formula[i].isdigit():
                i += 1
            return int(formula[j:i] or 1)

        def parse():
            nonlocal i
            cnt = collections.Counter()
            while i < n and formula[i] != ')':
                if formula[i] == '(':
                    i += 1
                    sub = parse()
                    i += 1
                    mult = read_num()
                    for name, c in sub.items():
                        cnt[name] += c * mult
                else:
                    j = i
                    i += 1
                    while i < n and formula[i].islower():
                        i += 1
                    name = formula[j:i]
                    cnt[name] += read_num()
            return cnt

        cnt = parse()
        return ''.join(name + (str(cnt[name]) if cnt[name] > 1 else '') for name in sorted(cnt))

解法 3：从右向左扫描 时间 O(n + m log m)，空间 O(n)

从右往左读化学式，先读到的数字是它左边原子或括号的倍数。用一个栈保存“当前累计倍数”：遇到右括号压入栈顶倍数乘以待用数字，遇到左括号弹出。遇到原子时把栈顶倍数乘以待用数字直接加到总计数里。

每个原子只累加一次，不需要逐层合并哈希表，所以与嵌套深度无关。反向读原子名和数字时边界容易写错。

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        cnt = collections.Counter()
        stack = [1]
        num = 1
        i = len(formula) - 1
        while i >= 0:
            c = formula[i]
            if c.isdigit():
                j = i
                while i >= 0 and formula[i].isdigit():
                    i -= 1
                num = int(formula[i + 1:j + 1])
            elif c == ')':
                stack.append(stack[-1] * num)
                num = 1
                i -= 1
            elif c == '(':
                stack.pop()
                i -= 1
            else:
                j = i
                while formula[i].islower():
                    i -= 1
                cnt[formula[i:j + 1]] += stack[-1] * num
                num = 1
                i -= 1
        return ''.join(name + (str(cnt[name]) if cnt[name] > 1 else '') for name in sorted(cnt))

栈与堆 专题的其他题目

查看 栈与堆 全部题目 →