Number of Atoms

hard stack string parsing hash map

Problem

Given a chemical formula string with element names, optional counts, and nested parenthesised groups, return the canonical formula listing each element followed by its total count (omit count of 1) in alphabetical order.

Inputformula="K4(ON(SO3)2)2"

Output"K4N2O14S4"

The inner SO3 is doubled, then the outer group is doubled again.

input

def countOfAtoms(formula):
    import re, collections
    tokens = re.findall(r'[A-Z][a-z]?|\d+|\(|\)', formula)
    stack = [collections.Counter()]
    i = 0
    while i < len(tokens):
        t = tokens[i]
        if t == '(': stack.append(collections.Counter()); i += 1
        elif t == ')':
            top = stack.pop(); i += 1
            k = int(tokens[i]) if i < len(tokens) and tokens[i].isdigit() else 1
            if k != 1 or True:
                if i < len(tokens) and tokens[i].isdigit(): i += 1
            for el, c in top.items(): stack[-1][el] += c * k
        else:
            el = t; i += 1
            k = int(tokens[i]) if i < len(tokens) and tokens[i].isdigit() else 1
            if i < len(tokens) and tokens[i].isdigit(): i += 1
            stack[-1][el] += k
    return ''.join(el + (str(stack[0][el]) if stack[0][el] > 1 else '') for el in sorted(stack[0]))

function countOfAtoms(formula) {
  const tokens = formula.match(/[A-Z][a-z]?|\d+|\(|\)/g);
  const stack = [new Map()];
  let i = 0;
  while (i < tokens.length) {
    const t = tokens[i];
    if (t === '(') { stack.push(new Map()); i++; }
    else if (t === ')') {
      const top = stack.pop(); i++;
      let k = 1;
      if (i < tokens.length && /^\d+$/.test(tokens[i])) { k = +tokens[i]; i++; }
      const cur = stack[stack.length - 1];
      for (const [e, c] of top) cur.set(e, (cur.get(e) || 0) + c * k);
    } else {
      const el = t; i++;
      let k = 1;
      if (i < tokens.length && /^\d+$/.test(tokens[i])) { k = +tokens[i]; i++; }
      const cur = stack[stack.length - 1];
      cur.set(el, (cur.get(el) || 0) + k);
    }
  }
  return [...stack[0]].sort().map(([e, c]) => e + (c > 1 ? c : '')).join('');
}

class Solution {
  public String countOfAtoms(String formula) {
    Deque<Map<String,Integer>> stack = new ArrayDeque<>();
    stack.push(new TreeMap<>());
    int i = 0, n = formula.length();
    while (i < n) {
      char ch = formula.charAt(i);
      if (ch == '(') { stack.push(new TreeMap<>()); i++; }
      else if (ch == ')') {
        Map<String,Integer> top = stack.pop(); i++;
        int j = i; while (j < n && Character.isDigit(formula.charAt(j))) j++;
        int k = j == i ? 1 : Integer.parseInt(formula.substring(i, j)); i = j;
        for (var e : top.entrySet()) stack.peek().merge(e.getKey(), e.getValue() * k, Integer::sum);
      } else {
        int j = i + 1; while (j < n && Character.isLowerCase(formula.charAt(j))) j++;
        String el = formula.substring(i, j); i = j;
        while (j < n && Character.isDigit(formula.charAt(j))) j++;
        int k = j == i ? 1 : Integer.parseInt(formula.substring(i, j)); i = j;
        stack.peek().merge(el, k, Integer::sum);
      }
    }
    StringBuilder sb = new StringBuilder();
    for (var e : new TreeMap<>(stack.peek()).entrySet()) { sb.append(e.getKey()); if (e.getValue() > 1) sb.append(e.getValue()); }
    return sb.toString();
  }
}

class Solution {
public:
  string countOfAtoms(string formula) {
    vector<map<string,int>> stk(1);
    int i = 0, n = formula.size();
    while (i < n) {
      char ch = formula[i];
      if (ch == '(') { stk.push_back({}); i++; }
      else if (ch == ')') {
        auto top = stk.back(); stk.pop_back(); i++;
        int j = i; while (j < n && isdigit(formula[j])) j++;
        int k = j == i ? 1 : stoi(formula.substr(i, j - i)); i = j;
        for (auto& [e, c] : top) stk.back()[e] += c * k;
      } else {
        int j = i + 1; while (j < n && islower(formula[j])) j++;
        string el = formula.substr(i, j - i); i = j;
        while (j < n && isdigit(formula[j])) j++;
        int k = j == i ? 1 : stoi(formula.substr(i, j - i)); i = j;
        stk.back()[el] += k;
      }
    }
    string res;
    for (auto& [e, c] : stk[0]) { res += e; if (c > 1) res += to_string(c); }
    return res;
  }
};

Explanation

A chemical formula has nested groups that can be multiplied, so we use a stack of counters where each level is a small map of element to count. The bottom of the stack accumulates the final totals.

First we tokenize the formula into elements (like SO), numbers, and parentheses. We start with one empty counter on the stack. When we read an element, we add its count to the counter at the top of the stack.

When we see ( we push a brand new counter to start a fresh group. When we see ) we pop that group's counter, read the multiplier k that follows (default 1), and merge every element back into the parent counter after multiplying each by k.

Example: "K4(ON(SO3)2)2". Inside, (SO3)2 doubles to S2 O6; the surrounding group (ON(SO3)2) then holds O7 N1 S2, and its outer 2 doubles everything to O14 N2 S4, plus the K4 at the base. Sorting the bottom counter gives "K4N2O14S4".

It works because each ) folds a completed group into its parent with the correct multiplier, so nesting is handled level by level, and a final alphabetical sort produces the canonical output.

Time: O(N^2) Space: O(N)