Problem statement

https://leetcode.com/problems/number-of-atoms/

Solution 1

Straightforward (but not easy to code solution) is the following: let helper(form) function return counter for all elements in form. Then we need to consider two cases:

  1. If we do not have ( inside our string, then we need to find first element: first letter will always be capital, then we traverse until we meet only lower letter. Then we find frequency of this element, which is either >1 or if it is empty, than it is 1. Finally, we recurrently call our function from the rest of the string.
  2. In the second case, we find the index of first (, and then find corresponding closing bracket. We run our helper 3 times: for part between first (, after corresponding it ) and between brackets. For part inside brackets we multiply it by frequency we get after ).

Complexity

Time complexity is potentially O(n^2), but it is already pretty good for this problem.

Code

class Solution:
    def countOfAtoms(self, formula):
        def findElem(elem, form, i):
            while i+1 < len(form) and form[i+1].islower():
                elem += form[i+1]
                i += 1
            return elem, i
        
        def findRep(rep, form, i):
            while i+1 < len(form) and form[i+1].isdigit():
                rep += form[i+1]
                i += 1
            rep = 1 if rep == "" else int(rep)
            return rep, i
        
        def helper(form):
            if not form: return Counter()
            if "(" not in form:
                elem, i = findElem(form[0], form, 0)
                rep, i = findRep("",form, i)
                return helper(form[i+1:]) + Counter({elem: rep})
            else:
                beg = form.index("(")
                balance = 1
                for i in range(beg+1, len(form)):
                    if form[i] == "(": balance += 1
                    if form[i] == ")": balance -= 1
                    if balance == 0: break
                inside = helper(form[beg+1:i])
                rep, i = findRep("",form, i)
                for k in inside: inside[k] = inside[k] * rep 
                return helper(form[:beg]) + helper(form[i+1:]) + inside
            
        t = sorted(helper(formula).items())
        return "".join(i + str(j)*(j!=1) for i,j in t)

Remark

Note, that we can make this solution O(n log n) or ever O(n) if we precalculate for each open brackets its corresponding closing bracket and if for each place we calculate the next closest place of open bracket.

Solution 2

There is another O(n) solution for this problem! The idea is to traverse our string from the end, and keep stack of frequencies. The trick, that if we do it in this way, we do not need to update all frequencies in stack, only the last one. num is current number we build and elem is current name of chemical element we built (in reversed order). Count is counter where we put final frequencies: we update them when we meet capital letter — it means that everything after this moment is processed.

When we meet ), it means that we need update stack: if we have num > 1, it means, that our frequency of element is increased: because we are one level deeper. If we meet (, we remove last element from stack, we are one level less deeper.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def countOfAtoms(self, formula):
        stack, num, elem, Count = [1], "", "", Counter()
        for s in formula[::-1]:
            if s.isdigit():
                num += s
            elif s == ')':
                k = int(num[::-1]) * stack[-1] if num else 1
                stack.append(k)
                num = ''
            elif s == '(':
                stack.pop()
                num = ''
            elif s.isupper():
                elem += s
                k = int(num[::-1]) if num else 1
                Count[elem[::-1]] += stack[-1] * k
                elem, num = "", ""
            elif s.islower():
                elem += s
            
        return "".join(i + str(j)*(j!=1) for i,j in sorted(Count.items()))