726. Number of Atoms

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  string countOfAtoms(string formula) {
    string ans;
    int i = 0;

    for (const auto& [elem, freq] : parse(formula, i)) {
      ans += elem;
      if (freq > 1)
        ans += to_string(freq);
    }

    return ans;
  }

 private:
  map<string, int> parse(const string& s, int& i) {
    map<string, int> count;

    while (i < s.length())
      if (s[i] == '(') {
        for (const auto& [elem, freq] : parse(s, ++i))
          count[elem] += freq;
      } else if (s[i] == ')') {
        const int num = getNum(s, ++i);
        for (auto&& [_, freq] : count)
          freq *= num;
        return count;  // Returns back to previous scope
      } else {         // s[i] must be uppercased
        const string& elem = getElem(s, i);
        const int num = getNum(s, i);
        count[elem] += num;
      }

    return count;
  }

  string getElem(const string& s, int& i) {
    const int elemStart = i++;  // s[elemStart] is uppercased
    while (i < s.length() && islower(s[i]))
      ++i;
    return s.substr(elemStart, i - elemStart);
  }

  int getNum(const string& s, int& i) {
    const int numStart = i;
    while (i < s.length() && isdigit(s[i]))
      ++i;
    const string& numString = s.substr(numStart, i - numStart);
    return numString.empty() ? 1 : stoi(numString);
  }
};

class Solution {
  public String countOfAtoms(String formula) {
    StringBuilder sb = new StringBuilder();
    Map<String, Integer> count = parse(formula);

    for (final String elem : count.keySet())
      sb.append(elem + (count.get(elem) == 1 ? "" : String.valueOf(count.get(elem))));

    return sb.toString();
  }

  private int i = 0;

  private Map<String, Integer> parse(String s) {
    Map<String, Integer> count = new TreeMap<>();

    while (i < s.length())
      if (s.charAt(i) == '(') {
        ++i; // Skip '('
        for (Map.Entry<String, Integer> entry : parse(s).entrySet()) {
          final String elem = entry.getKey();
          final int freq = entry.getValue();
          count.merge(elem, freq, Integer::sum);
        }
      } else if (s.charAt(i) == ')') {
        ++i; // Skip ')'
        final int num = getNum(s);
        for (final String elem : count.keySet()) {
          final int freq = count.get(elem);
          count.put(elem, freq * num);
        }
        return count; // Returns back to previous scope
      } else {
        final String elem = getElem(s);
        final int num = getNum(s);
        count.merge(elem, num, Integer::sum);
      }

    return count;
  }

  private String getElem(final String s) {
    final int elemStart = i++; // s[elemStart] is uppercased
    while (i < s.length() && Character.isLowerCase(s.charAt(i)))
      ++i;
    return s.substring(elemStart, i);
  }

  private int getNum(final String s) {
    final int numStart = i;
    while (i < s.length() && Character.isDigit(s.charAt(i)))
      ++i;
    final String numString = s.substring(numStart, i);
    return numString.isEmpty() ? 1 : Integer.parseInt(numString);
  }
}

class Solution:
  def countOfAtoms(self, formula: str) -> str:
    def parse() -> dict:
      ans = collections.defaultdict(int)

      nonlocal i
      while i < n:
        if formula[i] == '(':
          i += 1
          for elem, freq in parse().items():
            ans[elem] += freq
        elif formula[i] == ')':
          i += 1
          numStart = i
          while i < n and formula[i].isdigit():
            i += 1
          factor = int(formula[numStart:i])
          for elem, freq in ans.items():
            ans[elem] *= factor
          return ans
        elif formula[i].isupper():
          elemStart = i
          i += 1
          while i < n and formula[i].islower():
            i += 1
          elem = formula[elemStart:i]
          numStart = i
          while i < n and formula[i].isdigit():
            i += 1
          num = 1 if i == numStart else int(
              formula[numStart:i])
          ans[elem] += num

      return ans

    n = len(formula)

    ans = ""
    i = 0
    count = parse()

    for elem in sorted(count.keys()):
      ans += elem
      if count[elem] > 1:
        ans += str(count[elem])

    return ans