Number of Atoms

Number of Atoms is a classic parsing problem. It looks intimidating because of the nested parentheses, but it is really just a stack problem in disguise: every open paren starts a new tally, and every close paren folds that tally back into the one beneath it.

The slow way first

You might try to expand the formula textually — rewrite (OH)2 as OHOH, flatten every group, then count. But groups can nest inside groups, and a multiplier can be large, so expanding the string can blow up in size and is fiddly to get right. We want to count without ever materializing the expanded formula.

The question to ask: when I hit a close paren, how do I know which atoms belong to the group I am closing? I need a way to keep each group's running counts separate until I know its multiplier. That is exactly what a stack of count maps gives me.

The idea: a stack of count maps

Scan left to right, holding a stack where each entry is a count map for one open group. Start with one bottom map. For each character:

• A capital letter (plus any lowercase) is an element name; an optional number after it is its count. Add it to the top map.
• An open paren ( pushes a fresh empty map — new atoms now land in it.
• A close paren ) reads the number after it, multiplies the top map by that number, pops it, and merges those scaled counts into the map below.

The key insight: the top of the stack is always the map we add into. Parentheses just push and pop levels, and the multiply-merge on a close paren is what applies the group's exponent.

Walk through it

Step through the animation for Mg(OH)2. We add Mg into the bottom map L0. The ( pushes a fresh map L1. O and H land in L1. Then )2 multiplies L1 by 2 to get {O: 2, H: 2} and merges it into L0, leaving {Mg: 1, O: 2, H: 2}. Sorting the names gives H2MgO2.

Pseudocode

stack = [ empty map ]                 # one map per open group
i = 0
while i < length(formula):
    if formula[i] is a capital letter:
        read the full element name
        read the count after it (default 1)
        add count to the TOP map
    else if formula[i] == "(":
        push a fresh empty map; i += 1
    else:                              # ")"
        i += 1
        read the multiplier (default 1)
        pop the top map
        for each (name, c) in it:
            add c * multiplier into the new top map
return names sorted, each with its count if > 1

The Python solution

def count_of_atoms(formula):
    stack = [{}]
    i, n = 0, len(formula)
    while i < n:
        if formula[i].isupper():
            name, i = read_name(formula, i)
            cnt, i = read_int(formula, i, default=1)
            stack[-1][name] = stack[-1].get(name, 0) + cnt
        elif formula[i] == "(":
            stack.append({}); i += 1
        else:  # ")"
            i += 1
            mult, i = read_int(formula, i, default=1)
            top = stack.pop()
            for name, c in top.items():
                stack[-1][name] = stack[-1].get(name, 0) + c * mult
    return "".join(n + (str(c) if c > 1 else "") for n, c in sorted(stack[-1].items()))

• stack starts with one empty map; stack[-1] is always the map for the current group.
• read_name reads a capital plus any trailing lowercase letters; read_int reads consecutive digits, defaulting to 1 when none follow.
• A capital letter adds its count into the top map.
• ( pushes a fresh empty map so the group is tallied on its own.
• ) reads the multiplier, pops the top map, and merges each count c * mult into the map below — this is the multiply-merge that applies the group exponent.
• At the end the bottom map holds every atom; we sort the names and append a count only when it exceeds 1.

Complexity

We trade O(n) extra space (the stack of maps) for a clean single pass. The total cost is dominated by the scan plus a small sort over the distinct element names.

When this pattern shows up

Practice