leetcode/algorithms/python/numberofAtoms/numberofAtoms.py at main · MichelleZ/leetcode · GitHub

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
48
#! /usr/bin/env python3
# -*- coding: utf-8 -*-

# Source: https://leetcode.com/problems/number-of-atoms/
# Author: Miao Zhang
# Date:   2021-03-04

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        self.i = 0
        res = ''
        cnt = self.countAtoms(formula)
        for k in sorted(cnt.keys()):
            res += k
            if cnt[k] > 1:
                res += str(cnt[k])
        return res

    def countAtoms(self, formula: str) -> dict:
        counts = collections.defaultdict(int)
        while self.i != len(formula):
            if formula[self.i] == '(':
                self.i += 1
                tmp_counts = self.countAtoms(formula)
                count = self.getCount(formula)
                for k, v in tmp_counts.items():
                    counts[k] += v * count
            elif formula[self.i] == ')':
                self.i += 1
                return counts
            else:
                name = self.getName(formula)
                counts[name] += self.getCount(formula)
        return counts

    def getName(self, formula: str) -> str:
        name = ''
        while self.i < len(formula) and formula[self.i].isalpha() and (not name or formula[self.i].islower()):
            name += formula[self.i]
            self.i += 1
        return name

    def getCount(self, formula: str) -> int:
        cnt_str = ''
        while self.i < len(formula) and formula[self.i].isdigit():
            cnt_str += formula[self.i]
            self.i += 1
        return 1 if not cnt_str else int(cnt_str)