Description

Submission

class Solution {
public:
    string countOfAtoms(string formula) {
        int n = formula.size();
        map<string, int> cnt;
        stack<map<string, int>> stk;
        map<string, int> cur;

        for(int i = 0; i < n;) {
            auto ch = formula[i];

            if(ch == '(') {
                // push the previous one, start a new stack
                stk.push(cur);
                cur.clear(); // start a new one
                ++i;
            } else if(ch == ')') {
                // find the coefficient
                int j = i + 1;
                while(j < n && isdigit(formula[j])) ++j;
                int coeff = 1;
                if(j != i && isdigit(formula[j-1]))
                    coeff = stoi(formula.substr(i + 1, j - i - 1));
                for(auto& p: cur) {
                    p.second *= coeff;
                }

                // merge the previous and the current one
                // set the cur as the merged one
                for(auto [key, val]: stk.top()) {
                    cur[key] += val;
                }

                stk.pop();
                i = j;
            } else if(isupper(ch)) {
                    int j = i + 1;
                    while(j < n && islower(formula[j])) ++j;
                    string key = formula.substr(i, j - i);              
                    int val = 1;
                    int k = j;
                    while(k < n && isdigit(formula[k])) ++k;
                    if(j != k) {
                        val = stoi(formula.substr(j, k - j));
                    }
                    cur[key] += val;
                    i = k;
                }
        }

        stringstream ss;
        for(auto [key, val]: cur) {
            ss << key;
            if(val > 1) ss << val;
        }
        return ss.str();
    }
};

// charater types: upper, lower, digit, parenthesis

// K4(ON(SO3)2)2
// K4(ON..)2


// K4(ON(SO3)2(SO3)2)2

// 15