Python3 with Stack and defaultdict


  • 0
    S

    Idea is similar to many others, straightforward and simple

    class Solution:
        def countOfAtoms(self, formula):
            """
            :type formula: str
            :rtype: str
            """
            cur_ele = ''
            cur_num = 1
            stack = []
            cur_dict = collections.defaultdict(int)
            i = 0
            
            while (i < len(formula)):
                if formula[i].isupper():
                    if cur_ele is not '':
                        cur_dict[cur_ele] += cur_num
                    cur_num = 1
                    cur_ele = formula[i]
                elif formula[i].islower():
                    cur_ele += formula[i]
                elif formula[i].isdigit():
                    if formula[i - 1].isdigit():
                        cur_num = cur_num * 10 + int(formula[i])
                    else:
                        cur_num = int(formula[i])
                elif formula[i] == '(':
                    if cur_ele is not '':
                        cur_dict[cur_ele] += cur_num
                    cur_num = 1
                    cur_ele = ''
                    stack.append(cur_dict) 
                    cur_dict = collections.defaultdict(int)
                else: #formula[i] == ')'
                    if cur_ele != '':
                        cur_dict[cur_ele] += cur_num
                    i += 1
                    temp_count = 0
                    while i < len(formula) and formula[i].isdigit():
                        temp_count = temp_count * 10 + int(formula[i])
                        i += 1
                    if temp_count == 0:
                        temp_count = 1
                    i -=1
                    for ele in cur_dict.keys():
                        cur_dict[ele] = cur_dict[ele] * temp_count
                    prev_dict = stack.pop()
                    if len(prev_dict) != 0:
                        for ele in cur_dict.keys():
                            prev_dict[ele] += cur_dict[ele]
                        cur_dict = prev_dict
                    cur_num = 1
                    cur_ele = ''
                i += 1
                
            if cur_ele != '':
                cur_dict[cur_ele] += cur_num
            
            output = ''
            for k, v in sorted(cur_dict.items()):
                output += k
                if v != 1:
                    output += str(v)
            return output
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.