Among the solutions on the web, this is the most simple yet delightful code for solving leetcode 726-number-of-atoms I've ever seen.

726-number-of-atoms.py

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

        def _countOfAtoms(formula):
            nonlocal i
            counts = defaultdict(int)
            while i < len(formula):
                if formula[i] == "(":
                    i += 1
                    tmp_counts = _countOfAtoms(formula)
                    count = getCount(formula)
                    for k, v in tmp_counts.items():
                        counts[k] += v * count
                elif formula[i] == ")":
                    i += 1
                    return counts
                else:
                    name = getName(formula)
                    counts[name] += getCount(formula)
            return counts

        def getName(formula):
            nonlocal i
            name = ""
            while i < len(formula) and formula[i].isalpha() and (not name or formula[i].islower()):
                name += formula[i]
                i += 1
            return name

        def getCount(formula):
            nonlocal i
            count_str = ""
            while i < len(formula) and formula[i].isdigit():
                count_str += formula[i]
                i += 1
            return 1 if len(count_str) == 0 else int(count_str)

        return "".join(k if v == 1 else (k + str(v)) for k, v in sorted(_countOfAtoms(formula).items()))