diff options
Diffstat (limited to 'number_atoms.py')
-rw-r--r-- | number_atoms.py | 47 |
1 files changed, 47 insertions, 0 deletions
diff --git a/number_atoms.py b/number_atoms.py new file mode 100644 index 0000000..86fa2d7 --- /dev/null +++ b/number_atoms.py @@ -0,0 +1,47 @@ +# 726. Number of Atoms +from collections import defaultdict + +""" +given a string 'formula' representing a chemical formula, return the count of +each atom. the atomic element always stats with an uppercase character, then +zer or more lowercase letters, representing the name. one or more digits +representing that element's count may follow if the count is greater than 1. +if the count is 1, no digits will follow. return the count of all elements as +a string in the following form: the first name, followed by its count, +followed by the second name, followed by its count, and so on. +""" + + +class Solution(object): + def countOfAtoms(self, formula): + """ + :type formula: str + :rtype: str + """ + dic, coeff, stack, elem, cnt, i = defaultdict(int), 1, [], "", 0, 0 + for c in formula[::-1]: + if c.isdigit(): + cnt += int(c) * (10**i) + i += 1 + elif c == ")": + stack.append(cnt) + coeff *= cnt + i = cnt = 0 + elif c == "(": + coeff //= stack.pop() + i = cnt = 0 + elif c.isupper(): + elem += c + dic[elem[::-1]] += (cnt or 1) * coeff + elem = "" + i = cnt = 0 + elif c.islower(): + elem += c + return "".join(k + str(v > 1 and v or "") for k, v in sorted(dic.items())) + + +if __name__ == "__main__": + obj = Solution() + print(obj.countOfAtoms(formula="H2O")) + print(obj.countOfAtoms(formula="Mg(OH)2")) + print(obj.countOfAtoms(formula="K4(ON(SO3)2)2")) |