aboutsummaryrefslogtreecommitdiff
path: root/number_atoms.py
blob: 86fa2d7a27d63016764503b95dd98008ade24c0c (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
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"))