Want Help Cracking FAANG?

(Then click this)

 ×
Back to Question Bank

726. Number of Atoms - Leetcode Solution

Problem Description

The "Number of Atoms" problem on LeetCode asks you to parse a chemical formula string and count the number of each type of atom present in the formula. The formula may contain nested parentheses, and numbers (subscripts) that indicate the number of times an atom or group of atoms appears.

For example, given the input formula = "K4(ON(SO3)2)2", you must return a string like "K4N2O14S4" representing the count of each atom in lexicographical order.

  • Each atom is represented by a capital letter, possibly followed by lowercase letters (e.g., "H", "He").
  • Parentheses can be nested and are followed by a number (the multiplier for the group).
  • The output must list atoms in sorted order (lexicographically by atom name).
  • There is always one valid solution; the input is always a valid chemical formula.

Thought Process

At first, you might think to scan the string and count atoms as you see them. However, because of nested parentheses and multipliers, a simple scan won't work. For example, the atom count inside (SO3)2 must be multiplied by 2, and if that is nested deeper, it could be multiplied again by another number.

To solve this, we need a way to "remember" the multipliers for each group and apply them correctly, even when groups are nested. This naturally suggests using a stack to keep track of atom counts at each level of nesting. When we finish a group, we multiply all its counts by the group's multiplier and merge them with the outer level.

This is a classic example where a brute-force approach (just counting as you go) fails, and a structured approach using a stack is necessary for correctness and efficiency.

Solution Approach

The core idea is to use a stack to handle nested parentheses and multipliers. Here’s the step-by-step process:

  1. Initialization: Start with a stack containing an empty map (dictionary) to store atom counts at the current level.
  2. Parsing: Walk through the formula string character by character:
    • If you see '(', push a new empty map onto the stack.
    • If you see ')', pop the top map (this is the group just closed), parse the number after ')', multiply all atom counts in this map by that number, and merge them into the map below.
    • If you see an atom name, parse the full name (capital letter + possible lowercase letters), parse any number following it (default 1), and add/update the count in the current map on top of the stack.
  3. Finalization: After parsing, the bottom map in the stack contains the total atom counts. Sort the atom names and build the output string.
  4. Design Justification: Using a stack lets us handle arbitrarily deep nesting and ensures multipliers are applied in the correct order. A hash map (dictionary) is used for O(1) atom count lookups and updates.

Example Walkthrough

Let's trace the input "K4(ON(SO3)2)2" step by step:

  1. Start with a stack containing an empty map: [{{}}]
  2. Parse "K4": Add K:4 to the top map: [{{'K': 4}}]
  3. Encounter '(': Push a new map: [{{'K': 4}}, {{}}]
  4. Parse "O": Add O:1 to top: [{{'K': 4}}, {{'O': 1}}]
  5. Parse "N": Add N:1 to top: [{{'K': 4}}, {{'O': 1, 'N': 1}}]
  6. Encounter '(': Push a new map: [{{'K': 4}}, {{'O': 1, 'N': 1}}, {{}}]
  7. Parse "S": Add S:1: [{{'K': 4}}, {{'O': 1, 'N': 1}}, {{'S': 1}}]
  8. Parse "O3": Add O:3: [{{'K': 4}}, {{'O': 1, 'N': 1}}, {{'S': 1, 'O': 3}}]
  9. Encounter ')', followed by '2': Pop top map ({{'S': 1, 'O': 3}}), multiply by 2: {{'S': 2, 'O': 6}}, merge into previous: [{{'K': 4}}, {{'O': 1+6=7, 'N': 1, 'S': 2}}]
  10. Encounter ')', followed by '2': Pop top map ({{'O': 7, 'N': 1, 'S': 2}}), multiply by 2: {{'O': 14, 'N': 2, 'S': 4}}, merge into previous: [{{'K': 4, 'O': 14, 'N': 2, 'S': 4}}]
  11. End of string. Sort atoms and build output: "K4N2O14S4"

At each step, the stack structure ensures that multipliers are applied only to the correct group, even with deep or complex nesting.

Time and Space Complexity

  • Brute-force approach: Would require parsing and recalculating counts for every possible group, potentially O(N2) or worse, where N is the length of the formula.
  • Optimized stack-based approach:
    • Time Complexity: O(N), since each character is processed once, and map operations are O(1) on average.
    • Space Complexity: O(N), for the stack and the atom count maps, since in the worst case every character could be a parenthesis or atom.
  • The sorting step at the end is O(M log M), where M is the number of unique atoms, but M is typically much smaller than N.

Summary

The "Number of Atoms" problem teaches how to parse structured strings with nested groupings and multipliers. The key insight is to use a stack to manage nested contexts, ensuring multipliers are applied at the right scope. Using a hash map for atom counts allows efficient updates and merges. This approach is both efficient and elegant, handling complex nesting with clear, maintainable code.

Code Implementation

from collections import defaultdict

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        import re
        stack = [defaultdict(int)]
        i, n = 0, len(formula)
        while i < n:
            if formula[i] == '(':
                stack.append(defaultdict(int))
                i += 1
            elif formula[i] == ')':
                i += 1
                start = i
                while i < n and formula[i].isdigit():
                    i += 1
                mult = int(formula[start:i] or '1')
                popped = stack.pop()
                for atom, count in popped.items():
                    stack[-1][atom] += count * mult
            else:
                start = i
                i += 1
                while i < n and formula[i].islower():
                    i += 1
                atom = formula[start:i]
                start_num = i
                while i < n and formula[i].isdigit():
                    i += 1
                count = int(formula[start_num:i] or '1')
                stack[-1][atom] += count
        result = ''
        for atom in sorted(stack[-1]):
            cnt = stack[-1][atom]
            result += atom + (str(cnt) if cnt > 1 else '')
        return result
      
#include <string>
#include <stack>
#include <unordered_map>
#include <map>
using namespace std;

class Solution {
public:
    string countOfAtoms(string formula) {
        stack<unordered_map<string, int>> stk;
        stk.push(unordered_map<string, int>());
        int n = formula.size(), i = 0;
        while (i < n) {
            if (formula[i] == '(') {
                stk.push(unordered_map<string, int>());
                ++i;
            } else if (formula[i] == ')') {
                ++i;
                int start = i, mult = 1;
                while (i < n && isdigit(formula[i])) ++i;
                if (start < i) mult = stoi(formula.substr(start, i-start));
                auto top = stk.top(); stk.pop();
                for (auto &p : top) {
                    stk.top()[p.first] += p.second * mult;
                }
            } else {
                int start = i++;
                while (i < n && islower(formula[i])) ++i;
                string atom = formula.substr(start, i-start);
                int start_num = i, count = 1;
                while (i < n && isdigit(formula[i])) ++i;
                if (start_num < i) count = stoi(formula.substr(start_num, i-start_num));
                stk.top()[atom] += count;
            }
        }
        map<string, int> res;
        for (auto &p : stk.top()) res[p.first] = p.second;
        string ans;
        for (auto &p : res) {
            ans += p.first;
            if (p.second > 1) ans += to_string(p.second);
        }
        return ans;
    }
};
      
import java.util.*;
class Solution {
    public String countOfAtoms(String formula) {
        Deque<Map<String, Integer>> stack = new ArrayDeque<>();
        stack.push(new HashMap<>());
        int n = formula.length(), i = 0;
        while (i < n) {
            char c = formula.charAt(i);
            if (c == '(') {
                stack.push(new HashMap<>());
                i++;
            } else if (c == ')') {
                i++;
                int start = i, mult = 1;
                while (i < n && Character.isDigit(formula.charAt(i))) i++;
                if (start < i) mult = Integer.parseInt(formula.substring(start, i));
                Map<String, Integer> popped = stack.pop();
                for (String atom : popped.keySet()) {
                    stack.peek().put(atom, stack.peek().getOrDefault(atom, 0) + popped.get(atom) * mult);
                }
            } else {
                int start = i++;
                while (i < n && Character.isLowerCase(formula.charAt(i))) i++;
                String atom = formula.substring(start, i);
                int start_num = i, count = 1;
                while (i < n && Character.isDigit(formula.charAt(i))) i++;
                if (start_num < i) count = Integer.parseInt(formula.substring(start_num, i));
                stack.peek().put(atom, stack.peek().getOrDefault(atom, 0) + count);
            }
        }
        TreeMap<String, Integer> res = new TreeMap<>(stack.peek());
        StringBuilder sb = new StringBuilder();
        for (String atom : res.keySet()) {
            sb.append(atom);
            if (res.get(atom) > 1) sb.append(res.get(atom));
        }
        return sb.toString();
    }
}
      
var countOfAtoms = function(formula) {
    let stack = [new Map()];
    let i = 0, n = formula.length;
    while (i < n) {
        if (formula[i] === '(') {
            stack.push(new Map());
            i++;
        } else if (formula[i] === ')') {
            i++;
            let start = i;
            while (i < n && /\d/.test(formula[i])) i++;
            let mult = parseInt(formula.slice(start, i) || '1');
            let popped = stack.pop();
            for (let [atom, count] of popped.entries()) {
                stack[stack.length-1].set(atom, (stack[stack.length-1].get(atom) || 0) + count * mult);
            }
        } else {
            let start = i++;
            while (i < n && formula[i] >= 'a' && formula[i] <= 'z') i++;
            let atom = formula.slice(start, i);
            let startNum = i;
            while (i < n && /\d/.test(formula[i])) i++;
            let count = parseInt(formula.slice(startNum, i) || '1');
            stack[stack.length-1].set(atom, (stack[stack.length-1].get(atom) || 0) + count);
        }
    }
    let atoms = Array.from(stack[0].entries()).sort((a, b) => a[0].localeCompare(b[0]));
    let res = '';
    for (let [atom, count] of atoms) {
        res += atom + (count > 1 ? count : '');
    }
    return res;
};