Number of Atoms LeetCode solution in Python

1
class Solution:
2
    def countOfAtoms(self, formula: str) -> str:
3
        stack, elem, i = [{}], "", 0
4

5
        while i < len(formula):
6
            # Extract the full element name
7
            if formula[i].isupper():
8
                j = i + 1
9
                while j < len(formula) and formula[j].islower():
10
                    j += 1
11
                elem = formula[i:j]
12
                i = j
13
                # If no digits follow the element, assign a count of 1
14
                if (
15
                    i == len(formula)
16
                    or not formula[i].isdigit()
17
                    and not formula[i].islower()
18
                ):
19
                    stack[-1][elem] = stack[-1].get(elem, 0) + 1
20
            # Extract the count
21
            elif formula[i].isdigit():
22
                j = i
23
                while j < len(formula) and formula[j].isdigit():
24
                    j += 1
25
                count = int(formula[i:j])
26
                stack[-1][elem] = stack[-1].get(elem, 0) + count
27
                i = j
28
            # Handle open parentheses by pushing a new dict
29
            elif formula[i] == "(":
30
                stack.append({})
31
                i += 1
32
            # Handle close parentheses by merging with the previous dict
33
            elif formula[i] == ")":
34
                i += 1
35
                j = i
36
                while j < len(formula) and formula[j].isdigit():
37
                    j += 1
38
                multiplier = int(formula[i:j] or 1)
39
                top = stack.pop()
40
                for elem, count in top.items():
41
                    stack[-1][elem] = stack[-1].get(elem, 0) + count * multiplier
42
                i = j
43

44
        # Convert the result to the desired format
45
        atoms = sorted(stack[0].items())
46
        return "".join(
47
            [atom + (str(count) if count > 1 else "") for atom, count in atoms]
48
        )