aoc14.py

import copy 

class Group():
    def __init__(self, team, n_units, hp, weaknesses, immunity, damage, tdamage, initiative):
        self.team = team
        self.n_units = n_units
        self.hp = hp
        self.weaknesses = set(weaknesses)
        self.damage = damage
        self.tdamage = tdamage
        self.immunity = set(immunity)
        self.initiative = initiative
    
    def power(self):
        return self.n_units*self.damage
    
    def damaged_by(self, attacker):
        damage = attacker.power()
        if attacker.tdamage in self.immunity:
            return 0
        if attacker.tdamage in self.weaknesses:
            return (damage * 2)
        return damage
        
    def attacked_by(self, g):
        damage = self.damaged_by(g)
        self.n_units -= damage // self.hp

    def __repr__(self):
        return ("""TEAM {} has a group with {} each {} hp (immune to {}, weak to {}) attack {} {} at init {}\n""".format(self.team,
                                                                  self.n_units,
                                                                  self.hp,
                                                                  self.immunity,
                                                                  self.weaknesses,
                                                                  self.damage,
                                                                  self.tdamage,
                                                                  self.initiative
                                                                ))

def parse_props(props):
    immunities_on = False
    immunities = []
    weaknesses = []
    for word in props:
        if word == "immune":
            immunities_on = True
            continue
        if word == "to":
            continue
        if word == "weak":
            immunities_on = False
            continue
            
        if immunities_on:
            immunities.append(word.replace(",","").replace(";","").replace(" ",""))
        else:
            weaknesses.append(word.replace(",","").replace(";","").replace(" ",""))
            
    return immunities, weaknesses
    
    
def parse(team, lines):
    groups = []
    lines = lines.split("\n")[1:]
    immunities, weaknesses = [],[]
    for line in lines:
        parts = line.split("(")
        hp = parts[0].split()
        units = int(hp[0])
        hp = int(hp[4])
        try:
            props = parts[1].split(")")[0].split()
            immunities, weaknesses = parse_props(props)
            attack = parts[1].split(")")[1].split()
        except IndexError:
            immunities, weaknesses = [],[]
            attack = parts[0].split()[7:]
        
        tdamage = attack[6]
        initiative = int(attack[10])
        attack = int(attack[5])
      
        groups.append(Group(team, units, hp, weaknesses, immunities, attack, tdamage, initiative))
    return groups

def init(filename):
    with open(filename, "r") as file:
        data = file.read().split("\n\n")
        immune = parse("immune",data[0])
        infection = parse("infection",data[1])
        return immune, infection
        
def target_selection(immune, infection):
    opposition = {"immune":infection, "infection":immune}
    combats = {}
    unassigned = copy.copy(immune+infection)
    defenders = dict()
    while unassigned:
        attacker = max(unassigned, key=(lambda x: (x.power(),
                                                    x.initiative)))
        assert(attacker.n_units>0)
        unassigned.remove(attacker)
        # Target selection
        damages = dict()
        if(len(opposition[attacker.team])<1):
            continue
        #print(opposition[attacker.team])
        defender = max(opposition[attacker.team], key=(lambda x: (x.damaged_by(attacker),
                                                                  x.power(),
                                                                  x.initiative)))
        assert(defender.n_units>0)
        opposition[attacker.team].remove(defender)
        
        # defender chosen
        assert(attacker not in combats)
        assert(defender not in combats.values())
        combats[attacker] = defender
    return combats

def combat(immune, infection):
    opposition = {"immune":infection, "infection":immune}
    round_counter = -1
    while(len(immune)>0 and len(infection)>0):
        round_counter += 1
        combats = target_selection(copy.copy(immune), copy.copy(infection))
        while len(combats)>0:
            attacker = max(combats, key=(lambda x: x.initiative))
            defender = combats[attacker]
            del combats[attacker]
            if(attacker.n_units<1):
                print("dead")
                continue
            assert(defender.n_units>0)
            
            defender.attacked_by(attacker)
            
            if defender.n_units < 1:
                opposition[attacker.team].remove(defender)

    return immune, infection  

def part1(filename):
    immune, infection = init(filename)
    immune, infection = combat(immune, infection)
    return sum([x.n_units for x in immune+infection])

part1("24")