Given a chemical formula (given as a string), return the count of each atom.

An atomic element always starts with an uppercase character, then zero or more lowercase letters, representing the name.

1 or more digits representing the count of that element may follow if the count is greater than 1.  If the count is 1, no digits will follow.  For example, H2O and H2O2 are possible, but H1O2 is impossible.

Two formulas concatenated together produce another formula.  For example, H2O2He3Mg4 is also a formula.  

A formula placed in parentheses, and a count (optionally added) is also a formula.  For example, (H2O2) and (H2O2)3 are formulas.

Given a formula, output the count of all elements as a string in the following form: the first name (in sorted order), followed by its count (if that count is more than 1), followed by the second name (in sorted order), followed by its count (if that count is more than 1), and so on.

Example 1:

Input: 
formula = "H2O"
Output: "H2O"
Explanation: 
The count of elements are {'H': 2, 'O': 1}.



Example 2:

Input: 
formula = "Mg(OH)2"
Output: "H2MgO2"
Explanation: 
The count of elements are {'H': 2, 'Mg': 1, 'O': 2}.



Example 3:

Input: 
formula = "K4(ON(SO3)2)2"
Output: "K4N2O14S4"
Explanation: 
The count of elements are {'K': 4, 'N': 2, 'O': 14, 'S': 4}.



Note:
All atom names consist of lowercase letters, except for the first character which is uppercase.
The length of formula will be in the range [1, 1000].
formula will only consist of letters, digits, and round parentheses, and is a valid formula as defined in the problem.

Difficulty: interview

Starter code:

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        

Solutions:
class Solution(object):
     loc=0
     lastloc=-1
     f=''
     def getNext(self,formular,locked=False):
         stype=0 # 0:null, 1:numeric, 2/20: Elem, 3: parenthesis
         ret=0
         if self.loc==self.lastloc: return (0,0)
         i=self.loc
         while i <len(formular):
             if stype in (0,1) and  formular[i].isnumeric():
                 ret=int(formular[i])+ret*10
                 stype=1
             elif stype==0 and formular[i].isupper():
                 stype=20
                 ret=formular[i]
             elif stype in (20,2) and formular[i].islower():
                 stype=2
                 ret+=formular[i]
             elif stype==0 and formular[i] in "()":
                 stype=3+"()".index(formular[i])
                 ret=formular[i]
             else: break
             i+=1
         if not locked:
             self.lastloc=self.loc
             self.loc=i
         return (stype,ret)
     
     def countOfAtoms(self, formula):
         stk=[]
         cnt={}
         r=''
         self.loc=0
         n=self.getNext(formula)
         while n!=(0,0):
             if n[0] in (2,3,20):
                 stk.append([n[1],1])
             elif n[0]==1:
                 stk[-1][1]=n[1]
             elif n[0]==4:
                 time=1
                 i=-1
                 if self.getNext(formula,True)[0]==1:
                     time=self.getNext(formula)[1]
                 while stk[i][0]!='(':
                     stk[i][1]*=time
                     i-=1
                 stk[i][0]='$'
             n=self.getNext(formula)
         while any(stk):
             n=stk.pop()
             if n[0]!='$':
                 cnt[n[0]]=cnt.get(n[0],0)+n[1]
         for i in sorted(cnt.keys()):
             r+="%s%d"%(i,cnt[i]) if cnt[i]>1 else i
         return r

class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         stack = [collections.Counter()]
 
         i = 0
         length = len(formula)
         ss = ""
         while i < length:
             if formula[i].isupper():
                 ss = formula[i]
                 i += 1
                 while i < length:
                     if formula[i].islower():
                         ss += formula[i]
                         i += 1
                     else:
                         i -= 1
                         break
                 num = 0
                 i += 1
                 while i < length:
                     if formula[i].isdigit():
                         num = 10*num + int(formula[i])
                         i += 1
                     else:
                         i -= 1
                         break
                 num = num if num else 1
                 stack[-1][ss] = num if ss not in stack[-1] else stack[-1][ss]+num
 
             elif formula[i] == "(":
                 stack.append(collections.Counter())
             elif formula[i] == ")":
                 i += 1
                 multiplier = 0
                 while i < length:
                     if formula[i].isdigit():
                         multiplier = 10*multiplier + int(formula[i])
                         i += 1
                     else:
                         i -= 1
                         break
                 cur = stack.pop()
                 for n, v in list(cur.items()):
                     stack[-1][n] = stack[-1][n] + v*multiplier if n in stack[-1] else v*multiplier
 
             i += 1
         return "".join(name + (str(stack[-1][name]) if stack[-1][name] > 1 else '')
                                for name in sorted(stack[-1]))


class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         stack = []
         mapping = collections.defaultdict(int)
         s = ""
         num = 0
         i = 0
         while i < len(formula):
             c = formula[i]
             if c.islower():
                 s += c
                 i += 1
             elif c.isdigit():
                 num = 10 * num + int(c)
                 i += 1
             else:
                 if num == 0:
                     num = 1
                 if s != "":
                     mapping[s] += num
                 s = ""
                 num = 0
                 i += 1
                 if c.isupper():
                     s = c
                 elif c == '(':
                     stack.append(mapping)
                     mapping = collections.defaultdict(int)
                 elif c == ')':
                     while i < len(formula) and formula[i].isdigit():
                         num = 10 * num + int(formula[i])
                         i += 1
                     prev = stack.pop()
                     for k in mapping:
                         prev[k] += mapping[k] * num
                     mapping = prev
                     num = 0
         if num == 0:
             num = 1
         if s != "":
             mapping[s] += num
         ans = []
         for k, v in sorted(mapping.items()):
             ans.append(k)
             if v > 1:
                 ans.append(str(v))
         return ''.join(ans)


class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         dict = self.helper(formula)
         result = ""
         for key in sorted(dict):
             result += str(key) + (str(dict[key]) if dict[key] > 1 else "")
         return result
     
     def helper(self, formula):
         result = {}
         j = 0 # conventional i, j are reversed
         element = ""
         while j < len(formula):
             if formula[j] == "(":
                 if len(element) > 0:
                     result[element] = (result[element] if element in result else 0) + 1
                     element = ""
                 count = 1
                 end = j + 1
                 while end < len(formula) and count > 0:
                     if formula[end] == "(":
                         count += 1
                     if formula[end] == ")":
                         count -= 1
                     end += 1
                 dict = self.helper(formula[j + 1: end - 1])
                 i = end 
                 while i < len(formula) and formula[i].isdigit():
                     i += 1
                 m = int(formula[end: i])
                 for key, value in list(dict.items()):
                     result[key] = value * m + (result[key] if key in result else 0)
             elif formula[j].isdigit():
                 i = j + 1
                 while i < len(formula) and formula[i].isdigit():
                     i += 1
                 result[element] = (result[element] if element in result else 0) + int(formula[j:i])
                 element = ""
             else:
                 if len(element) > 0:
                     result[element] = (result[element] if element in result else 0) + 1
                     element = ""
                 i = j + 1
                 while i < len(formula) and formula[i].islower():
                     i += 1
                 element = formula[j:i]
             j = i
         if not formula[j - 1].isdigit():
             result[element] = (result[element] if element in result else 0) + 1
         return result
 


class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         formula+=' '
         stack,i=[{}],0
         while i<len(formula):
             if formula[i]>='A' and formula[i]<='Z':
                 j=i+1
                 while formula[j]>='a' and formula[j]<='z':
                     j+=1
                 k=j
                 while formula[k]>='0' and formula[k]<='9':
                     k+=1
                 t=1 if k==j else int(formula[j:k])
                 stack[-1][formula[i:j]]=stack[-1].get(formula[i:j],0)+t
                 i=k
             elif formula[i]=='(':
                 stack.append({})
                 i+=1
             elif formula[i]==')':
                 j=i+1
                 while formula[j]>='0' and formula[j]<='9':
                     j+=1
                 t=1 if j==i+1 else int(formula[i+1:j])
                 for x in stack[-1]:
                     stack[-2][x]=stack[-2].get(x,0)+stack[-1][x]*t
                 stack.pop()
                 i=j
             else:
                 i+=1
         a,s=sorted(stack[-1].keys()),''
         for x in a:
             s+=x if stack[-1][x]==1 else x+str(stack[-1][x])
         return s

class Solution:
     
     def __init__(self):
         self.i = 0
         self.formula = ''
     
     
     @staticmethod
     def _merge_into(dict_main, dict_added):
         """Change dict_main in place by adding the dict_added"""
         for key, val in list(dict_added.items()):
             dict_main[key] = dict_main.get(key, 0) + val
     
     @staticmethod
     def _multiply(dict_main, mult):
         """Return a new dict from the multiplication of everything by a number"""
         return {k: v * mult for k, v in list(dict_main.items())}    
     
     # Processing methods
     def cur(self):
         return self.formula[self.i]
     
     def consume(self):
         res = self.formula[self.i]
         self.i += 1
         return res
     
     def atend(self):
         return self.i == len(self.formula)
 
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         # Adding a space for easier test than always testing if we are at the end
         self.formula = formula
         
         res = {}
         while not self.atend():
             self._merge_into(res, self.process_next_part())
             
             
         res_str = []
         for key, value in sorted(res.items()):
             res_str.append(key)
             if value != 1:
                 res_str.append(str(value))
         return ''.join(res_str)
         
         
     def process_next_part(self):
         # Return i, dict_atoms
         cur = self.cur()
         if cur.isalpha():
             name = self._consume_name()
             value = self._consume_digit()
             return {name: value}
         elif cur == '(':
             local_res = {}
             self.consume()  # Consuming the '('
             while self.cur() != ')':
                 self._merge_into(local_res, self.process_next_part())
             self.consume()  # Consuming the ')'
             value = self._consume_digit()
             return self._multiply(local_res, value)
         else:
             raise Exception("char unkwon")
 
 
     def _consume_name(self):
         name = self.consume()
         if not self.atend() and self.cur().isalpha() and self.cur().lower() == self.cur():
             return name + self.consume()
         return name
         
     def _consume_digit(self):
         if self.atend() or not self.cur().isdigit():
             return 1
         res_str = self.consume()
         while not self.atend() and self.cur().isdigit():
             res_str += self.consume()
         return int(res_str)
 


class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         
         def merge(d1, d2):
             for k, v in list(d2.items()):
                 if k in d1:
                     d1[k] += v
                 else:
                     d1[k] = v
                     
         def mul(d, multiplier):
             for k in list(d.keys()):
                 d[k] *= multiplier
                 
         def helper(start):
             c = formula[start]
             els = {}
             end = start
             if c == '(':
                 end += 1
                 while formula[end] != ')':
                     end, els2 = helper(end)
                     merge(els, els2)
                 end += 1
             elif 'A' <= c <= 'Z':
                 end += 1
                 while end < len(formula) and 'a' <= formula[end] <= 'z':
                     end += 1
                 els[formula[start:end]] = 1
             else:
                 raise ValueError('Unexpected token')
                 
             end2 = end
             while end2 < len(formula) and '0' <= formula[end2] <= '9':
                 end2 += 1
             count = int(formula[end:end2]) if end2 != end else 1
             mul(els, count)
             return end2, els
         
         
         i, result = helper(0)
         while i < len(formula):
             i, els = helper(i)
             merge(result, els)
             
         keys = sorted(result.keys())
         s = ""
         for k in keys:
             if result[k] > 1:
                 s+="%s%d"%(k, result[k])
             else:
                 s+=k
         return s
         
             


from collections import Counter
 
 def tokenizer(formula):
     result = []
     
     current_token = ""
     
     numbers = "0123456789"
     
     def append_current_token():
         nonlocal current_token, result
         
         if current_token:
             result.append(current_token)
             current_token = ""
     
     for char in formula:
         if char == "(":
             append_current_token()
             result.append(char)
             
         elif char == ")":
             append_current_token()
             result.append(char)
         
         elif char in numbers:
             if current_token and current_token[-1].isalpha():
                 append_current_token()
             
             current_token += char
             
         elif char.isalpha():
             if current_token and (current_token[-1] in numbers or char.isupper()) :
                 append_current_token()
             
             current_token += char
             
     
     append_current_token()
     
     return result
     
 
 class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         result = Counter()
         formula = tokenizer(formula)
         
         mult_stack = []
         mult = 1
         
         precedent_number = 1
         
         for token in formula[::-1]:
             if token.isdigit():
                 n = int(token)
                 precedent_number = n
             
             elif token.isalpha():
                 result[token] += mult * precedent_number
                 precedent_number = 1
             
             elif token == ")":
                 mult_stack.append(precedent_number)
                 mult *= precedent_number
                 precedent_number = 1
             
             elif token == "(":
                 n = mult_stack[-1]
                 mult_stack.pop()
                 
                 mult //= n
         
         
         
         str_result = []
         
         for token in sorted(result.keys()):
                 str_result.append(token + str(result[token]) if result[token] != 1 else token)
         
         return ''.join(str_result)

class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         atoms = self.helper(formula)
         return ''.join([e + (str(c) if c > 1 else '') for e,c in sorted(atoms.items())])
     
     
     def helper(self, formula):
         i = 0
         atoms = {}
         while i < len(formula):
             if formula[i] == '(':
                 end = self.end_parens(formula, i)
                 atom = self.helper(formula[i+1:end])
             else:
                 end = self.end_element(formula, i)
                 atom = {formula[i:end+1]:1}
                 
             i = end + 1
             mult = 0
             while i < len(formula) and formula[i].isdigit():
                 mult = mult * 10 + int(formula[i])
                 i += 1
                 
             if mult == 0:
                 mult = 1
             
             for element, count in atom.items():
                 atoms.setdefault(element, 0)
                 atoms[element] += count * mult
         
         return atoms
                 
 
     def end_element(self, formula, i):
         i += 1
         while i < len(formula) and formula[i].islower():
             i += 1
             
         return i - 1
 
 
     def end_parens(self, formula, start):
         parens = 1
         i = start + 1
         while i < len(formula) and parens != 0:
             if formula[i] in '()':
                 parens += 1 if formula[i] == '(' else -1
             i += 1
         
         return i - 1

class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         stack = []
         mapping = collections.defaultdict(int)
         s = ""
         num = 0
         i = 0
         while i < len(formula):
             c = formula[i]
             if c.islower():
                 s += c
                 i += 1
             elif c.isdigit():
                 num = 10 * num + int(c)
                 i += 1
             else:
                 if s != "":
                     mapping[s] = mapping[s] + num if num > 0 else mapping[s] + 1
                 s = ""
                 num = 0
                 i += 1
                 if c.isupper():
                     s = c
                 elif c == '(':
                     stack.append(mapping)
                     mapping = collections.defaultdict(int)
                 elif c == ')':
                     while i < len(formula) and formula[i].isdigit():
                         num = 10 * num + int(formula[i])
                         i += 1
                     prev = stack.pop()
                     for k in mapping:
                         prev[k] += mapping[k] * num
                     mapping = prev
                     num = 0
         if s != "":
             mapping[s] = mapping[s] + num if num > 0 else mapping[s] + 1
         ans = []
         for k, v in sorted(mapping.items()):
             ans.append(k)
             if v > 1:
                 ans.append(str(v))
         return ''.join(ans)


class Solution:
     def countOfAtoms(self, formula):
         """
         :type formula: str
         :rtype: str
         """
         
         # stack = [dict()]
         stack = [collections.Counter()]
         i = 0
         while i < len(formula):
             if formula[i] == '(':
                 # stack.append(dict())
                 stack.append(collections.Counter())
                 print('new stack', stack)
                 i += 1
             elif formula[i] == ')':
                 i += 1
                 
                 num_start = i
                 
                 while i < len(formula) and formula[i].isdigit():
                     i += 1
                 
                 mult = formula[num_start:i]
                 
                 if mult is '':
                     mult = 1
                 else:
                     mult = int(formula[num_start:i])
                     
                 top = stack[-1]
                 print('top: ', top)
                 print('new mult: ', mult)
                 stack.pop()
                 print('poped top stack: ', stack)
                 for name in top:
                     stack[-1][name] = stack[-1][name] + top[name]*mult
             else:
                 name = formula[i]
                 i += 1
                 while i < len(formula) and formula[i].islower():
                     name += formula[i]
                     i += 1
                     
                 print('name of atom: ', name)
                     
                 num_start = i
                 
                 while i < len(formula) and formula[i].isdigit():
                     i += 1
                 
                 mult = formula[num_start:i]
                 
                 if mult is '':
                     mult = 1
                 else:
                     mult = int(formula[num_start:i])
                 
                 print('mult of {}: '.format(name), mult)
                 
                 stack[-1][name] += mult
                 print('updated current stack: ', stack)
         
         return "".join(name + (str(stack[-1][name]) if stack[-1][name] > 1 else '')
                        for name in sorted(stack[-1]))