726. Number of Atoms

Problem

Question

Given a string formula representing a chemical formula, return the count of each atom.

The atomic element always starts with an uppercase character, then zero or more lowercase letters, representing the name.

One or more digits representing that element’s count may follow if the count is greater than 1. If the count is 1, no digits will follow.

• For example, "H2O" and "H2O2" are possible, but "H1O2" is impossible.

Two formulas are concatenated together to produce another formula.

• For example, "H2O2He3Mg4" is also a formula.

A formula placed in parentheses, and a count (optionally added) is also a formula.

• For example, "(H2O2)" and "(H2O2)3" are formulas.

Return the count of all elements as a string in the following form: the first name (in sorted order), followed by its count (if that count is more than 1), followed by the second name (in sorted order), followed by its count (if that count is more than 1), and so on.

The test cases are generated so that all the values in the output fit in a 32-bit integer.

Example 1:

Input: formula = “H2O” Output: “H2O” Explanation: The count of elements are {‘H’: 2, ‘O’: 1}.

Example 2:

Input: formula = “Mg(OH)2” Output: “H2MgO2” Explanation: The count of elements are {‘H’: 2, ‘Mg’: 1, ‘O’: 2}.

Example 3:

Input: formula = “K4(ON(SO3)2)2” Output: “K4N2O14S4” Explanation: The count of elements are {‘K’: 4, ‘N’: 2, ‘O’: 14, ‘S’: 4}.

Constraints:

• 1 <= formula.length <= 1000
• formula consists of English letters, digits, '(', and ')'.
• formula is always valid.

---

Solutions

1. Reverse Traversal

Time Complexity: $O(n)$ | Space Complexity: $O(n)$

To solve this problem, we can iterate from the end of the given string while keeping track of the total multiplier. To store the counts, we use a dictionary. To manage multipliers within parentheses, a stack is used. Generally, there are four conditions to handle:

• When a digit is observed:

  • Continue iterating to build the entire number.
  • Save the number in cur_mult.

• When a ”)” is observed:

  • Push cur_mult onto the stack.
  • Update total_mult by multiplying it by cur_mult.
  • Reset cur_mult to 1.

• When a ”(” is observed:

  • Pop the stack.
  • Use the popped value to update total_mult by dividing it by this value.

• When a character is observed:

  • Continue iterating to find the entire chemical element (until the character is not lowercase).
  • Update the dictionary counter by adding the result of multiplying cur_mult with total_mult.
  • Reset cur_mult.`

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        
        stack = []
        total_mult = 1 
        cur_mult = 1
        memo = defaultdict(int)
 
        i = len(formula) - 1
        while i >= 0:
            cur = formula[i]
            if cur == ")":
                stack.append(cur_mult)
                total_mult *= cur_mult
                cur_mult = 1
            elif cur == "(":
                old_mult = stack.pop()
                total_mult /= old_mult
            elif cur.isdigit():
                number = 0
                multiplier = 1
                while i >= 0 and formula[i].isdigit():
                    number += multiplier * int(formula[i])
                    multiplier *= 10
                    i -= 1
                cur_mult = number if number != 0 else 1
                continue
            else:
                if cur.islower():
                    while i >= 0 and formula[i].islower():
                        cur = formula[i - 1] + cur
                        i -= 1
                memo[cur] += (total_mult * cur_mult)
                cur_mult = 1
            i -= 1
 
        keys = list(memo)
        keys.sort()
 
        result = ''
        for key in keys:
            number = str(int(memo[key])) if memo[key] > 1 else ''
            result = result + key + number
 
        return result
 

---

2. Recursive Approach

Time Complexity: $O(n)$ | Space Complexity: $O(n)$

In this method, most of the logic is similar, but we iterate forward. We call a helper function each time we encounter new parentheses. The recursive function returns a dictionary containing all the inner counts. We then multiply all these values by the number that follows the parentheses, effectively distributing the multiplier to all nested elements.

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        
        def iterate(fr): 
            memo = defaultdict(int)
            n = len(formula)
            while fr < n:
                cur = formula[fr]
 
                if cur == "(":
                    sub_memo, fr = iterate(fr + 1)
                    number = 0
                    while fr < n and formula[fr].isdigit():
                        number = number * 10 + int(formula[fr])
                        fr += 1
                    if number == 0: number = 1
                    for key in sub_memo:
                        memo[key] += sub_memo[key] * number
                elif cur == ")":
                    return memo, fr + 1
                else:
                    if cur.isalpha() and cur.isupper():
                        key = cur
                        fr += 1
                        while fr < n and formula[fr].isalpha() and formula[fr].islower():
                            key += formula[fr]
                            fr += 1
                        number = 0
                        while fr < n and formula[fr].isdigit():
                            number = number * 10 + int(formula[fr])
                            fr += 1
                        if number == 0: number = 1
                        memo[key] += number
 
            return memo, n
 
        memo, _ = iterate(0)
 
        result = []
        for atom in sorted(memo.keys()):
            count = memo[atom]
            if count == 1:
                result.append(atom)
            else:
                result.append(atom + str(count))
        
        return ''.join(result)