leetcode-go/leetcode/hard/726.cc

#include <map>
#include <string>
#include <vector>

#include <iostream>

using std::map;
using std::string;
using std::vector;

using std::cout; using std::endl;

class Solution {
public:
    string countOfAtoms(string formula) {
        map<string, int> atoms = parse_formula(formula);
        string result = "";
        for (auto& it : atoms) {
            result += it.first;
            if (1 < it.second)
                result += std::to_string(it.second);
        }
        return result;
    }
private:
    enum Symbol {UPPER, LOWER, DIGIT, LPARENT, RPARENT, END, ILLEGAL};

    struct Token {
        string tok;
        bool simple;
        int cnt;
    };

    Symbol symbol_of(char ch) {
        if ('A' <= ch && ch <= 'Z')
            return UPPER;
        if ('a' <= ch && ch <= 'z')
            return LOWER;
        if ('0' <= ch && ch <= '9')
            return DIGIT;
        if (ch == '(')
            return LPARENT;
        if (ch == ')')
            return RPARENT;
        return ILLEGAL;
    }

    vector<Token> tokenize(string formula) {
        vector<Token> result;
        int n = formula.size();
        for (int i = 0; i < n;) {
            switch (symbol_of(formula[i])) {
                case UPPER: {
                    string tok = read_tok(formula, i);
                    int cnt = read_cnt(formula, i);
                    result.push_back(Token{tok, true, cnt});
                    break;
                }
                case LPARENT: {
                    string tok = read_formula(formula, i);
                    int cnt = read_cnt(formula, i);
                    result.push_back(Token{tok, false, cnt});
                    break;
                }
                default: // Dead loop if invalid
                    break;
            }
        }
        return result;
    }

    string read_tok(string& str, int& i) {
        int beg = i++;
        if (i != str.size() && symbol_of(str[i]) == LOWER) {
            i++;
            return str.substr(beg, 2);
        }
        return str.substr(beg, 1);
    }

    int read_cnt(string& str, int& i) {
        int n = str.size();
        if (i == n || symbol_of(str[i]) != DIGIT)
            return 1;
        int beg = i;
        for (; i < n && symbol_of(str[i]) == DIGIT; i++);
        return std::stoi(str.substr(beg, i - beg));
    }

    string read_formula(string& str, int& i) {
        int n = str.size();
        int beg = ++i;
        int stack = 1;
        for (; i < n && 0 < stack; i++) {
            switch (symbol_of(str[i])) {
                case LPARENT:
                    stack++;
                    break;
                case RPARENT:
                    stack--;
                    break;
                default:
                    break;
            }
        }
        string formula = str.substr(beg, i - 1 - beg);
        return formula;
    }

    map<string, int> merge(vector<map<string, int>> forms) {
        map<string, int> result;
        for (auto& it : forms)
            merge(it, result);
        return result;
    }

    map<string, int> merge(map<string, int>& src, map<string, int>& dst) {
        for (auto& it : src)
            dst[it.first] = dst[it.first] + it.second;
        return dst;
    }

    map<string, int> parse_formula(string& formula) {
        vector<Token> tokens = tokenize(formula);
        vector<map<string, int>> forms;
        map<string, int> simple;
        for (auto& it : tokens) {
            if (it.simple) {
                simple[it.tok] = simple[it.tok] + it.cnt;
                continue;
            }
            map<string, int> comp = parse_formula(it.tok);
            for (auto& t : comp)
                t.second *= it.cnt;
            forms.push_back(comp);
        }
        forms.push_back(simple);
        return merge(forms);
    }

    void print_tokens(vector<Token>& tokens) {
        for (auto& it : tokens)
            cout << "{tok:" << it.tok << " sim:" << it.simple << " cnt:" << it.cnt << "}; ";
        cout << endl;
    }
};

int main() {
    Solution* sol = new Solution;

    string formula = "H2O";
    cout << sol->countOfAtoms(formula) << endl;

    formula = "H2O2((C(H2O)H)12)10";
    cout << sol->countOfAtoms(formula) << endl;

    formula = "K4(ON(SO3)2)2";
    cout << sol->countOfAtoms(formula) << endl;

    return 0;
}