Number of atoms

XXX. Number of Atoms

Stack

HashMap

String Parsing

Problem Statement:

Given a chemical formula as a string (e.g., "Mg(OH)2", "K4(ON(SO3)2)2"), return the count of each atom in the formula in lexicographical order.

Algorithm:

1. Use a stack to handle nested parentheses. Each level of parentheses has its own map to count atoms.
2. Traverse the formula character by character:
• If the character is '(', push the current map onto the stack and start a new level.
• If the character is ')', parse the multiplier after the closing parenthesis, multiply the counts in the current map, and merge them with the parent map from the stack.
• If the character is an uppercase letter, parse the full atom name and its count, updating the current map.
3. At the end, merge the counts from all levels into a global map.
4. Sort the keys of the global map lexicographically and construct the result string.

Complexity:

Time Complexity: O(n + k log k), where n is the length of the formula and k is the number of unique atoms.
Space Complexity: O(n + k) for the stack and maps.

Java Implementation:

public class Solution {
    public String countOfAtoms(String formula) {
        // This map stores the final atom counts
        Map globalMap = new HashMap<>();
        // Parse the formula and populate the global map
        parseFormula(formula, globalMap);

        // Sort the map keys (atoms) to construct the result in lexicographical order
        StringBuilder result = new StringBuilder();
        for (String atom : new TreeMap<>(globalMap).keySet()) {
            result.append(atom);
            if (globalMap.get(atom) > 1) 
                result.append(globalMap.get(atom));
        }
        return result.toString();
    }

    private void parseFormula(String formula, Map globalMap) {
        // Stack to store maps for handling nested parentheses
        Deque> stack = new ArrayDeque<>();
        // Current map to count atoms at the current level
        Map currentMap = new HashMap<>();
        int n = formula.length();
        int i = 0;

        while (i < n) {
            char c = formula.charAt(i);

            if (c == '(') {
                // Push the current map onto the stack and start a new level
                stack.push(currentMap);
                currentMap = new HashMap<>();
                i++;
            } else if (c == ')') {
                // Parse the multiplier after the ')'
                i++;
                int start = i;
                while (i < n && Character.isDigit(formula.charAt(i))) {
                    i++;
                }
                int multiplier = start < i ? Integer.parseInt(formula.substring(start, i)) : 1;

                // Multiply the counts in the current map and merge with the map from the stack
                Map parentMap = stack.pop();
                for (String atom : currentMap.keySet()) {
                    parentMap.put(atom, parentMap.getOrDefault(atom, 0) + currentMap.get(atom) * multiplier);
                }
                currentMap = parentMap; // Return to the previous level
            } else if (Character.isUpperCase(c)) {
                // Parse an atom name
                int start = i;
                i++;
                while (i < n && Character.isLowerCase(formula.charAt(i))) {
                    i++;
                }
                String atom = formula.substring(start, i);

                // Parse the count (if any)
                start = i;
                while (i < n && Character.isDigit(formula.charAt(i))) i++;
                
                int count = start < i ? Integer.parseInt(formula.substring(start, i)) : 1;
                // Update the count for the atom in the current map
                currentMap.put(atom, currentMap.getOrDefault(atom, 0) + count);
            }
        }

        // At the end, merge all counts into the global map
        for (String atom : currentMap.keySet()) 
            globalMap.put(atom, globalMap.getOrDefault(atom, 0) + currentMap.get(atom));
    }
}
                

Python Implementation:

def count_of_atoms(formula: str) -> str:
    from collections import defaultdict
    import re

    def parse_formula(formula, global_map):
        stack = []  # Stack to store maps for nested parentheses
        current_map = defaultdict(int)  # Current map at this level
        i, n = 0, len(formula)

        while i < n:
            if formula[i] == '(':
                stack.append(current_map)
                current_map = defaultdict(int)
                i += 1
            elif formula[i] == ')':
                i += 1
                start = i
                while i < n and formula[i].isdigit():
                    i += 1
                multiplier = int(formula[start:i]) if start < i else 1
                parent_map = stack.pop()
                for atom, count in current_map.items():
                    parent_map[atom] += count * multiplier
                current_map = parent_map
            elif formula[i].isupper():
                start = i
                i += 1
                while i < n and formula[i].islower():
                    i += 1
                atom = formula[start:i]
                start = i
                while i < n and formula[i].isdigit():
                    i += 1
                count = int(formula[start:i]) if start < i else 1
                current_map[atom] += count
        for atom, count in current_map.items():
            global_map[atom] += count

    global_map = defaultdict(int)
    parse_formula(formula, global_map)
    return ''.join(f"{atom}{global_map[atom] if global_map[atom] > 1 else ''}" for atom in sorted(global_map))
                

C++ Implementation:

#include 
#include 
#include 
#include