317-counting-elements.py

# from https://www.reddit.com/r/dailyprogrammer/comments/6eerfk/20170531_challenge_317_intermediate_counting/

import re

challenges = """CCl2F2
NaHCO3
C4H8(OH)2
PbCl(NH3)2(COOH)2"""
stack = [{
    'elements': {},
    'prev_element': ''
}]

pop_next = False


def handle_pop(multiplier):
    global pop_next

    if (pop_next):
        pop_next = False

        popped = stack.pop()
        for key, val in popped['elements'].items():
            add(stack[-1]['elements'], key, val * multiplier)


def add(elements, piece, amount=1):
    if elements.get(piece):
        elements[piece] += amount
    else:
        elements[piece] = amount


def handle(kind, piece):
    global stack
    global pop_next

    elements = stack[-1]['elements']
    prev_element = stack[-1]['prev_element']

    if kind == 'number':
        handle_pop(int(piece))
        elements[prev_element] += (int(piece) - 1)
    else:
        handle_pop(1)

    if kind == 'element':
        stack[-1]['prev_element'] = piece
        add(elements, piece)

    if kind == 'open':
        stack.append({
            'elements': {},
            'prev_element': ''
        })

    if kind == 'close':
        pop_next = True


def match(piece):
    tokens = [
        ('[A-Z]([a-z])?', 'element'),
        ('([0-9])+', 'number'),
        ('\(', 'open'),
        ('\)', 'close')
    ]

    for reg, kind in tokens:
        if re.compile(reg).match(piece):
            return handle(kind, piece)

    raise RuntimeError('Could not match %s to any type' % piece)


def solve(challenge):
    reg = re.compile('([A-Z][a-z]?|[0-9]+|\(|\))')
    pieces = reg.findall(challenge)

    for piece in pieces:
        match(piece)


def solve_all(challenges):
    for challenge in challenges.split('\n'):
        result = solve(challenge)

        print(stack[0]['elements'])


solve_all(challenges)