advent-of-code/2018/day-15/02.py

from pprint import pprint
from time import sleep

m = list()
hitpoints = dict()


with open('input.txt', 'r') as f:
    for raw_line in f.readlines():
        line = list()
        for field in raw_line.strip():
            line.append(field)
        m.append(line)

for x in range(len(m)):
        for y in range(len(m[x])):
            if m[x][y] in ['G', 'E']:
                hitpoints[(x, y)] = 200

def get_patch(pos):
    x, y = pos
    if x < 0 or x >= len(m):
        return None
    if y < 0 or y >= len(m[x]):
        return None
    return m[x][y]

def neighbors(pos):
    return [
        (pos[0], pos[1] + 1),
        (pos[0], pos[1] - 1),
        (pos[0] + 1, pos[1]),
        (pos[0] - 1, pos[1])
    ]

def valid_neighbors(pos):
    return [x for x in neighbors(pos) if get_patch(x) == '.']
    

def flood(pos):
    patches = set([pos])
    flood_map = {
        pos: 0
    }
    new = set(valid_neighbors(pos))
    i = 1
    while new:
        patches.update(new)
        new_new = set()
        for p in new:
            flood_map[p] = i
            new_new.update(valid_neighbors(p))
        new = new_new - patches
        i += 1
    return flood_map
    
def get_units():
    """unit is a tuple (pos, symbol)"""
    units = list()
    for x in range(len(m)):
        for y in range(len(m[x])):
            if m[x][y] in ['G', 'E']:
                units.append(((x, y), m[x][y]))
    return units

def get_closest_target(flood_map, targets):
    reachable_targets = {pos: dist for pos, dist in flood_map.items() if pos in targets}
    if len(reachable_targets) == 0:
        return None
    # print('Reachable targets:', reachable_targets)
    min_range = min(reachable_targets.values())
    closest = [pos for pos, dist in reachable_targets.items() if dist == min_range]
    # print('Closest targets:', closest)
    closest.sort()
    return closest[0]

def get_enemies(unit):
    return [e for e in get_units() if e[1] != unit[1]]

def get_targets(unit):
    targets = set()
    for e in get_enemies(unit):
        targets.update(valid_neighbors(e[0]))
    return targets

def get_enemies_in_range(unit):
    x, y = unit[0]
    enemies = get_enemies(unit)
    neighbors = [
        (x + 1, y),
        (x - 1, y),
        (x, y + 1),
        (x, y - 1)
    ]
    return [e for e in enemies if e[0] in neighbors]

def move(unit, target):
    move_to = dict()
    neighbors = valid_neighbors(unit[0])
    # print('neighbors', neighbors)
    for n in neighbors:
        flood_map = flood(n)
        # print(n, target, flood_map.get('target'))
        if target in flood_map:
            move_to[n] = flood_map[target]
    # print(move_to)
    min_range = min(move_to.values())
    closest = [pos for pos, dist in move_to.items() if dist == min_range]
    closest.sort()

    x1, y1 = unit[0]
    x2, y2 = closest[0]
    m[x1][y1] = '.'
    m[x2][y2] = unit[1]

    htp = hitpoints[unit[0]]
    del hitpoints[unit[0]]
    hitpoints[closest[0]] = htp

    return (closest[0])

def count_dwarfes():
    c = 0
    for line in m:
        for sym in line:
            if sym == 'E':
                c += 1
    return c

def count_goblins():
    c = 0
    for line in m:
        for sym in line:
            if sym == 'G':
                c += 1
    return c

def print_map():
    print('  ', ''.join([ str(y // 10) for y, _ in enumerate(m[0]) ]))
    print('  ', ''.join([ str(y % 10) for y, _ in enumerate(m[0]) ]))
    for i, line in enumerate(m):
        htp_list = [(pos, htp) for pos, htp in hitpoints.items() if pos[0] == i]
        htp_list.sort()
        print(str(i).ljust(2), ''.join(line), [h[1] for h in htp_list])
    print('  ', ''.join([ str(y % 10) for y, _ in enumerate(m[0]) ]))
    print('  ', ''.join([ str(y // 10) for y, _ in enumerate(m[0]) ]))


def main(elf_atk=3):
    i = 0
    while 1:
        # print()
        # print(f'After {i} rounds:')
        # print_map()
        units = get_units()
        units.sort()
        for u in units:
            if get_patch(u[0]) != u[1]:
                continue
            if count_dwarfes() == 0 or count_goblins() == 0:
                return i
            # print(f'UNIT: {u[1]} at {u[0]}')
            e = get_enemies_in_range(u)
            if len(e) == 0:
                # move
                targets = get_targets(u)
                closest = get_closest_target(flood(u[0]), targets)
                if closest is None:
                    # print('do nothing')
                    continue
                d = move(u, closest)
                u = (d, u[1])
                # print(f'Moved towards {closest} via {d}')
            e = get_enemies_in_range(u)
            if len(e) != 0:
                # attack
                atk_pos = [en[0] for en in e]
                htp_map = {pos: val for pos, val in hitpoints.items() if pos in atk_pos}
                min_htp = min(htp_map.values())
                atk_targets = [pos for pos, htp in htp_map.items() if htp == min_htp]

                target = atk_targets[0]
                # print(f'Unit {u} attacking {target}')
                atk = elf_atk if u[1] == 'E' else 3
                hitpoints[target] -= atk
                if hitpoints[target] <= 0:
                    del hitpoints[target]
                    x, y = target
                    m[x][y] = '.'
                
        i += 1
        # pprint(hitpoints)
        # sleep(.1)
    
print('Elves before:', count_dwarfes())
i = main(15)
print('Elves after:', count_dwarfes())
# print()
# print(f'After {i + 1} rounds:')
# print_map()    
print()
print('Number of complete rounds:', i)
print('hitpoints:', sum(hitpoints.values()))
print('Result', i * sum(hitpoints.values()))