726. 原子的数量

给你一个字符串化学式 formula ，返回 每种原子的数量 。

原子总是以一个大写字母开始，接着跟随 0 个或任意个小写字母，表示原子的名字。

如果数量大于 1，原子后会跟着数字表示原子的数量。如果数量等于 1 则不会跟数字。

例如，"H2O" 和 "H2O2" 是可行的，但 "H1O2" 这个表达是不可行的。
两个化学式连在一起可以构成新的化学式。

例如 "H2O2He3Mg4" 也是化学式。
由括号括起的化学式并佐以数字（可选择性添加）也是化学式。

例如 "(H2O2)" 和 "(H2O2)3" 是化学式。
返回所有原子的数量，格式为：第一个（按字典序）原子的名字，跟着它的数量（如果数量大于 1），然后是第二个原子的名字（按字典序），跟着它的数量（如果数量大于 1），以此类推。

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/number-of-atoms
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

---

import java.util.*;

public class Solution {
    private int index;

    private String str;

    private void merge(Map<String, Integer> m1, Map<String, Integer> m2) {
        for (Map.Entry<String, Integer> entry: m2.entrySet()) {
            int count = m1.getOrDefault(entry.getKey(), 0);
            count += entry.getValue();
            m1.put(entry.getKey(), count);
        }
    }

    private void multiply(Map<String, Integer> m, int num) {
        for (Map.Entry<String, Integer> entry: m.entrySet()) {
            int count = m.getOrDefault(entry.getKey(), 0);
            count *= num;
            m.put(entry.getKey(), count);
        }
    }

    private static void add(Map<String, Integer> m, String str, int num) {
        m.put(str, m.getOrDefault(str, 0) + num);
    }

    public String countOfAtoms(String formula) {
        Map<String, Integer> countMap = solve(formula);
        StringBuilder sb = new StringBuilder();
        for (Map.Entry<String, Integer> entry: countMap.entrySet()) {
            sb.append(entry.getKey());
            if  (entry.getValue() != 1) {
                sb.append(entry.getValue());
            }
        }
        return sb.toString();
    }

    public TreeMap solve(String str) {
        this.index = 0;
        this.str = str;
        Stack<Map<String, Integer>> stack = new Stack<>();
        stack.push(new HashMap<>());
        while (index < str.length()) {
            if (str.charAt(index) == '(') {
                stack.push(new HashMap<>());
                index++;
            } else if (str.charAt(index) == ')') {
                index++;
                int num = getNextInt();
                Map<String, Integer> m1 = stack.pop();
                Map<String, Integer> m2 = stack.peek();
                multiply(m1, num);
                merge(m2, m1);
            } else {
                add(stack.peek(), getNextString(), getNextInt());
            }
        }
        return new TreeMap<>(stack.pop());
    }

    private int getNextInt() {
        if (index == str.length() || ! Character.isDigit(str.charAt(index))) {
            return 1;
        }
        int num = 0;
        while (index < str.length() && Character.isDigit(str.charAt(index))) {
            num = num * 10 + str.charAt(index++) - '0';
        }
        return num;
    }

    private String getNextString() {
        StringBuilder sb = new StringBuilder();
        sb.append(str.charAt(index++));
        while(index < str.length() && Character.isLowerCase(str.charAt(index))) {
            sb.append(str.charAt(index++));
        }
        return sb.toString();
    }
}