种种类型题目利用栈实现很方便。
第一次实现的时候超时了,原因在于利用一个带有重复元素的列表来保存当前遇到的原子。
这样的话列表可能非常大,效率极低。利用字典解决了这个问题。
from collections import defaultdict
class Solution:
def countOfAtoms(self, formula: str) -> str:
i = 0
len1 = len(formula)
tem_dic = defaultdict(int)
stack = []
set1 = set('ABCDEFGHIGKLMNOPQRSTUVWXYZ')
set2 = set('abcdefghigklmnopqrstuvwxyz')
set3 = set('0123456789')
set4 = set()
while i < len1:
if formula[i] in set1:
c = formula[i]
i += 1
while i < len1 and formula[i] in set2:
c += formula[i]
i += 1
set4.add(c)
if i < len1 and formula[i] in set3:
val = int(formula[i])
i += 1
while i < len1 and formula[i] in set3:
val = val * 10 + int(formula[i])
i += 1
tem_dic[c] += val
else:
tem_dic[c] += 1
elif formula[i] == '(':
stack.append(tem_dic)
tem_dic = defaultdict(int)
i += 1
else:
i += 1
if i < len1 and formula[i] in set3:
val = int(formula[i])
i += 1
while i < len1 and formula[i] in set3:
val = val * 10 + int(formula[i])
i += 1
for key in tem_dic:
tem_dic[key] = tem_dic[key] * val
res_dic = stack.pop()
for key in res_dic:
tem_dic[key] = tem_dic[key] + res_dic[key]
res = []
list1 = list(set4)
list1.sort()
for val in list1:
if tem_dic[val] > 1:
res.append(val + str(tem_dic[val]))
else:
res.append(val)
return ''.join(res)