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

## C++

``````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;  // return 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);
}
};
``````

## JAVA

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

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.put(elem, count.getOrDefault(elem, 0) + freq);
}
} 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; // return back to previous scope
} else {
final String elem = getElem(s);
final int num = getNum(s);
count.put(elem, count.getOrDefault(elem, 0) + num);
}

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);
}
}
``````

## Python

``````class Solution:
def countOfAtoms(self, formula: str) -> str:
def parse() -> dict:
ans = 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
``````