Neat Python with Explanation - 35ms

  • 6

    The first regex ([A-Z]{1}[a-z]?|\(|\)|\d+) splits up the input string into a few kinds of tokens for parsing: (1) An atom (2) A number (3) Open bracket (4) Closing bracket. These are the only tokens we need to do our parsing.

    An input of Mg(OH)2 will be tokenized into: ['Mg', '(', 'O', 'H', ')', '2'].
    An input of K4(ON(SO3)2)2 will be tokenized into: ['K', '4', '(', 'O', 'N', '(', 'S', 'O', '3', ')', '2', ')', '2'].

    As we iterate through the tokens, there are three cases that we need to handle:

    1. Open bracket - We push a new dictionary onto a stack to keep track of the atoms and its count in this current group
    2. Close bracket - The next token might be a number/count. Check whether if it is a count. If it is, multiply all the atoms at the top of the stack by the count and combine it with a dictionary below it in the stack.
    3. Normal atom - The next token might be a number/count. Check whether if it is a count. If it is, add that atom and its count to the top of the stack.

    Cases 2 and 3 are very similar, so we can combine them.

    At the end, sort the atoms alphabetically and format them nicely to be returned.

    - Yangshun

    class Solution(object):
        def countOfAtoms(self, formula):
            import re
            from collections import defaultdict
            tokens = list(filter(lambda c: c, re.split('([A-Z]{1}[a-z]?|\(|\)|\d+)', formula)))
            stack, i = [defaultdict(int)], 0
            while i < len(tokens):
                token = tokens[i]
                if token == '(':
                    count = 1
                    # Check if next token is a number.
                    if i + 1 < len(tokens) and'^\d+$', tokens[i + 1]):
                        count, i = int(tokens[i + 1]), i + 1
                    atoms = stack.pop() if token == ')' else { token: 1 }
                    # Combine counts of atoms.
                    for atom in atoms:
                        stack[-1][atom] += atoms[atom] * count
                i += 1
            return ''.join([atom + (str(count) if count > 1 else '') for atom, count in sorted(stack[-1].items())])

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.