Number of Atoms — Solution

Problem

Given a chemical formula string containing element names, counts, and parentheses, count the total number of each atom and return them formatted in alphabetical order. An element name starts with an uppercase letter and may be followed by lowercase letters. A count of 1 is omitted from the output. Parentheses followed by a number multiply all enclosed atom counts by that number.

For example, "Mg(OH)2" expands to one Mg, two O, and two H atoms.

• Input: formula = "Mg(OH)2"
• Output: "H2MgO2"
• Explanation: The group (OH) is doubled by the 2, giving O×2 and H×2; combined with the leading Mg, output is sorted alphabetically.

Another example shows deeper nesting:

• Input: formula = "K4(ON(SO3)2)2"
• Output: "K4N2O14S4"
• Explanation: Inner (SO3)2 → S×2, O×6; outer (ONS2O6)2 → O×2, N×2, S×4, O×12; plus K×4 → K4N2O14S4.

Intuition

Parentheses in a formula create a nested scope — the multiplier after ) applies to everything accumulated inside. A stack of frequency maps handles this naturally: push a fresh counter when entering (, then pop and scale when leaving ), merging results back into the enclosing scope. Everything outside any group goes directly into the current top-of-stack counter.

Solution — Stack Parser

Scan left to right maintaining a stack of counters. On ( push a new scope; on ) scale the popped scope and merge it into the one below; on an uppercase letter parse the full element name and its optional count into the current scope.

1. Initialize a stack with one empty frequency map representing the outermost scope.
2. When encountering (, push a new empty map and advance.
3. When encountering ), advance past it, then read the multiplier (default 1 if no digits follow). Pop the top map, multiply every count by the multiplier, and add each scaled count to the new top map.
4. When encountering an uppercase letter, read the full element name (uppercase + following lowercase letters), then read the integer count (default 1 if no digits follow). Increment that element's entry in the top map.
5. After scanning the full formula, one map remains on the stack. Sort its keys alphabetically and build the result string, appending the count only when it exceeds 1.

1def count_of_atoms(formula: str) -> str:
2    stack = [{}]
3    i, n = 0, len(formula)
4
5    while i < n:
6        if formula[i] == '(':
7            stack.append({})
8            i += 1
9        elif formula[i] == ')':
10            i += 1
11            start = i
12            while i < n and formula[i].isdigit():
13                i += 1
14            multiplier = int(formula[start:i]) if start < i else 1
15            top = stack.pop()
16            for element, count in top.items():
17                # scale this group's counts before merging into enclosing scope
18                stack[-1][element] = stack[-1].get(element, 0) + count * multiplier
19        else:  # uppercase letter — start of element name
20            start = i
21            i += 1
22            while i < n and formula[i].islower():
23                i += 1
24            element = formula[start:i]
25            start = i
26            while i < n and formula[i].isdigit():
27                i += 1
28            count = int(formula[start:i]) if start < i else 1
29            stack[-1][element] = stack[-1].get(element, 0) + count
30
31    result = []
32    for element in sorted(stack[0]):
33        result.append(element)
34        if stack[0][element] > 1:
35            result.append(str(stack[0][element]))
36    return "".join(result)

Time: O(n²) worst case — each ) merges the popped map into the one below, and with d levels of nesting containing k elements each, merge work totals O(k × d); deeply nested formulas like ((A)2)2... hit O(n²).

Space: O(n) — the stack holds at most n/2 maps, and total entries across all maps are bounded by the formula length.

Complexity Summary

Common Mistakes

• Forgetting the default count of 1 — an element or group without a following number has an implicit count of 1; formula[start:i] will be empty when no digits are present, so the int(...) conversion must be guarded with if start < i else 1.
• Not resetting start before the count parsing loop — the start variable is reused: once for the element name slice and once for the digit slice. Forgetting to update start to i after parsing the element name means the count extraction reads letters as digits and crashes.
• Using a single flat dict instead of a stack — a flat dict cannot handle nested groups where the multiplier outside ) should only apply to the atoms inside those parentheses, not to everything accumulated so far.
• Omitting the count when it equals 1 — the output format requires "H2O" not "H2O1"; the count should only be appended when it strictly exceeds 1.
• Building the result with an unsorted map — Python's dict and C++'s unordered_map do not sort by key; the final output must explicitly sort element names with sorted(...) (Python) or map<string,int> (C++) to guarantee alphabetical order.

Related Problems

• decode-string — same stack pattern with brackets and multipliers, applied to character strings instead of element counts
• evaluate-reverse-polish-notation — stack-based expression evaluation where each operator pops and combines values
• valid-parentheses — simpler stack problem that checks bracket matching without multipliers or nested counters
• parsing-a-boolean-expression — recursive expression parsing with operators and nested sub-expressions
• longest-valid-parentheses — stack tracks bracket positions to find the longest valid bracketed substring