Day 19: Aplenty

# Advent of Code 2023
# Day 19: Aplenty

from copy import deepcopy
from functools import reduce
from json import loads
from operator import mul
from re import sub

workflows = {}
parts = []

with open("inputs/day_19.txt") as file:
    half = 0

    for line in file:
        if half == 0:
            if line == "\n":
                half = 1
            else:
                halves = line.split("{")
                halves[1] = halves[1].rstrip("\n")[:-1]

                conditions = []

                for condition in halves[1].split(","):
                    if condition == "A":
                        conditions.append(True)
                    elif condition == "R":
                        conditions.append(False)
                    elif ":" not in condition:
                        conditions.append(condition)
                    else:
                        condition_halves = condition.split(":")

                        if condition_halves[1] == "A":
                            result = True
                        elif condition_halves[1] == "R":
                            result = False
                        else:
                            result = condition_halves[1]

                        conditions.append((
                            condition_halves[0][1],
                            "xmas".index(condition_halves[0][0]),
                            int(condition_halves[0][2:]),
                            result
                        ))

                workflows[halves[0]] = conditions
        else:
            parts.append(list(loads(sub(
                r"([xmas])=", r'"\g<1>": ', line.rstrip("\n")
            )).values()))

rating_sum = 0

for part in parts:
    workflow = workflows["in"]
    done = False

    while not done:
        for condition in workflow:
            if condition == True:
                rating_sum += sum(part)
                done = True
                break
            elif condition == False:
                done = True
                break
            elif isinstance(condition, str):
                workflow = workflows[condition]
            else:
                if condition[0] == "<":
                    pass_ = part[condition[1]] < condition[2]
                else:
                    pass_ = part[condition[1]] > condition[2]

                if pass_:
                    if condition[3] == True:
                        rating_sum += sum(part)
                        done = True
                    elif condition[3] == False:
                        done = True
                    else:
                        workflow = workflows[condition[3]]

                    break

print(f"Part 1: {rating_sum}")

states = (["in", [1, 4000], [1, 4000], [1, 4000], [1, 4000]],)
combinations = 0

while states:
    new_states = []

    for state in states:
        for condition in workflows[state[0]]:
            if condition == True:
                combinations += reduce(
                    mul, (range_[1] - range_[0] + 1 for range_ in state[1:])
                )

                break
            elif isinstance(condition, str):
                state[0] = condition
                new_states.append(state)
                break
            elif condition != False:
                state_pass = deepcopy(state)
                range_pass = state_pass[condition[1] + 1]

                range_pass[condition[0] == "<"] = (
                    condition[2] + (1, -1)[condition[0] == "<"]
                )

                state[condition[1] + 1][condition[0] == ">"] = condition[2]

                if range_pass[0] <= range_pass[1]:
                    if condition[3] == True:
                        combinations += reduce(mul, (
                            range_[1] - range_[0] + 1
                            for range_ in state_pass[1:]
                        ))
                    elif condition[3] != False:
                        state_pass[0] = condition[3]
                        new_states.append(state_pass)

    states = new_states

print(f"Part 2: {combinations}")

Today is the day I by far did the most involved custom data structures for, I had an Instruction class and an Item class and a RequiredState class. For a while, RequiredState was a list of Requirements for each letter.

Finally (after way too long) I realized that each RequiredState is going to be a contiguous range. Upon realizing that, this got WAY easier. I started tracking like this instead:

class RequiredState:
    def __init__(self):
        self.min_values = {
            'x': 1,
            'm': 1,
            'a': 1,
            's': 1,
        }
        self.max_values = {
            'x': 4000,
            'm': 4000,
            'a': 4000,
            's': 4000,
        }

and updating like so:

    def add_req(self, kind: str, op, value, invert_this):
        if invert_this:
            op = '>' if op == '<' else '<'
        else:
            # change to strict equality
            value = value + (-1 if op == '<' else 1)
        if op == '<':
            self.max_values[kind] = min(self.max_values[kind], value)
        else:
            self.min_values[kind] = max(self.min_values[kind], value)

pub fn part2(input: PuzzleInput) -> isize {
    let mut flows = HashMap::default();
    let mut accepted = Vec::new();

    let workflows = Rc::new(input.workflows);

    let flow = Flow {
        part: PartRange {
            x: 1..4001,
            m: 1..4001,
            a: 1..4001,
            s: 1..4001,
        },
        current_workflow: "in".to_string(),
        current_index: 0,
        workflows: workflows.clone(),
        accepted: 4000 * 4000 * 4000 * 4000,
    };

    flows.insert(flow, 1);

    loop {
        let mut new_flows = utility_belt::misc::state_iteration(&flows, transition, ());

        for (flow, count) in new_flows.iter_mut() {
            if flow.current_workflow == "A" {
                accepted.push(flow.clone());
            }
        }

        new_flows.retain(|f, _| f.current_workflow != "A" && f.current_workflow != "R");

        if new_flows.is_empty() {
            break;
        }

        flows = new_flows;
    }

    accepted.iter().map(|x| x.accepted).sum::<usize>() as isize
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct PartRange {
    pub x: Range<usize>,
    pub m: Range<usize>,
    pub a: Range<usize>,
    pub s: Range<usize>,
}

#[derive(Clone, Debug)]
pub struct Flow {
    pub part: PartRange,
    pub current_workflow: String,
    pub current_index: usize,
    pub workflows: Rc<BTreeMap<String, Workflow>>,
    pub accepted: usize,
}

impl Hash for Flow {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.part.hash(state);
        self.current_workflow.hash(state);
        self.current_index.hash(state);
    }
}

impl PartialEq for Flow {
    fn eq(&self, other: &Self) -> bool {
        self.part == other.part
            && self.current_workflow == other.current_workflow
            && self.current_index == other.current_index
    }
}

impl Eq for Flow {}

fn transition(from: &Flow, input: &()) -> Vec<Flow> {
    let mut result = Vec::default();

    let workflow = from.workflows.get(&from.current_workflow).unwrap();
    let rule = workflow.rules.get(from.current_index);

    if rule.is_none() {
        let mut new_flow = from.clone();
        new_flow.current_workflow = workflow.default.clone();
        new_flow.current_index = 0;
        result.push(new_flow);
        return result;
    }

    let rule = rule.unwrap();

    let ((lower_next, lower_next_index), (upper_next, upper_next_index)) = match rule.comparison {
        Comparison::LessThan => (
            (rule.next.clone(), 0),
            (from.current_workflow.clone(), from.current_index + 1),
        ),
        Comparison::GreaterThan => (
            (from.current_workflow.clone(), from.current_index + 1),
            (rule.next.clone(), 0),
        ),
    };

    match rule.category {
        Category::X => {
            let full_count = from.part.x.end - from.part.x.start;

            let (lower_range, upper_range) = match rule.comparison {
                Comparison::LessThan => {
                    (from.part.x.start..rule.value, rule.value..from.part.x.end)
                }

                Comparison::GreaterThan => (
                    from.part.x.start..(rule.value + 1),
                    (rule.value + 1)..from.part.x.end,
                ),
            };

            if !lower_range.is_empty() {
                let lower_count = lower_range.end - lower_range.start;
                result.push(Flow {
                    part: PartRange {
                        x: lower_range,
                        m: from.part.m.clone(),
                        a: from.part.a.clone(),
                        s: from.part.s.clone(),
                    },
                    current_workflow: lower_next,
                    current_index: lower_next_index,
                    accepted: from.accepted * lower_count / full_count,
                    ..from.clone()
                });
            };

            if !upper_range.is_empty() {
                let upper_count = upper_range.end - upper_range.start;
                result.push(Flow {
                    part: PartRange {
                        x: upper_range,
                        m: from.part.m.clone(),
                        a: from.part.a.clone(),
                        s: from.part.s.clone(),
                    },
                    current_workflow: upper_next,
                    current_index: upper_next_index,
                    accepted: from.accepted * upper_count / full_count,
                    ..from.clone()
                });
            };
        }

        Category::M => {
            [...]
         }
   
        Category::A => {
            [...]
        }
     
        Category::S => {
            [...]
         }
    }

    result
}

The only slightly notable thing about this is that I abstracted the "Iterate a state transition function" into a re-usable library function after day 12. I didn't think I'd need it so soon again.

def build_workflow(raw_filters: str, default: str) -> Callable[[Part], str]:
    def _get_next_destination(part: Part):
        for raw_filter in raw_filters.split(","):
            category, op, value, destination = extract_filter_components(raw_filter)
            if op(part[category], value):
                return destination

        return default

    return _get_next_destination

def build_counting_workflow(
    raw_filters: str, default: str
) -> Callable[[RangedPart], list[tuple[str, RangedPart]]]:
    def _get_next_destinations(part: RangedPart):
        ranges: list[tuple[str, RangedPart]] = []

        for raw_filter in raw_filters.split(","):
            category, op, value, dest = extract_filter_components(raw_filter)

            if op == gt:
                keep, send = bisect_range(part[category], value + 1)
            else:
                send, keep = bisect_range(part[category], value)

            ranges.append((dest, {**part, category: send}))
            part = {**part, category: keep}

        # whatever is left also goes
        return ranges + [(default, part)]

    return _get_next_destinations

module Part1 (solve) where

import AOCTools (splitBy, splitOn)
import qualified Data.Map as Map
import Data.Map ((!))

solve :: String -> Int
solve str = sum (map partRating (filter (testPart workflows) parts)) where
    (workflows, parts) = parseInput str

type ID = String

data Workflow = Workflow { name :: ID, rules :: [Rule], final :: ID }

type Rule = (Part -> Bool, ID)

data Part = Part { x :: Int, m :: Int, a :: Int, s :: Int }

parseRule :: String -> Rule
parseRule str = (pred, res) where
    (predStr, res) = splitOn ":" str
    letter : op : val = predStr
    v = read val
    o = case op of
        '>' -> (>)
        '<' -> (<)
    l = case letter of
        'x' -> x
        'm' -> m
        'a' -> a
        's' -> s
    pred p = o (l p) v

parseWorkflow :: String -> Workflow
parseWorkflow str = Workflow name rules final where
    name = takeWhile (/= '{') str
    innerStr = takeWhile (/='}') (drop (length name + 1) str)
    pieces = splitBy "," innerStr
    rules = map parseRule (take (length pieces - 1) pieces)
    final = pieces !! (length pieces - 1)

partRating :: Part -> Int
partRating (Part xVal mVal aVal sVal) = xVal + mVal + aVal + sVal

parsePart :: String -> Part
parsePart str = Part xVal mVal aVal sVal where
    pieces = splitBy "," (takeWhile (/='}') (tail str))
    [xVal, mVal, aVal, sVal] = map (read . drop 2) pieces

parseInput :: String -> ([Workflow], [Part])
parseInput s = (workflows, parts) where 
    (workflowStrs, partsStrs) = splitOn "\n\n" s
    workflows = map parseWorkflow $ lines workflowStrs
    parts = map parsePart $ lines partsStrs

testPart :: [Workflow] -> Part -> Bool
testPart workflows part = testPartHelper "in" where
    workflowMap = Map.fromList (map (\w -> (name w, w)) workflows)
    testPartHelper :: ID -> Bool
    testPartHelper "A" = True
    testPartHelper "R" = False
    testPartHelper idNum = testPartHelper $ processPart (workflowMap ! idNum) part

processPart :: Workflow -> Part -> ID
processPart workflow part = helper (rules workflow) where
    helper [] = final workflow
    helper ((rule, result):xs) = if rule part then result else helper xs

module Part2 (solve) where

import AOCTools (splitBy, splitOn)
import qualified Data.Map as Map
import Data.Map (Map, (!))

solve :: String -> Integer
solve str = result where
    workflows = parseInput str
    workflowMap = Map.fromList (map (\w -> (name w, w)) workflows)
    validParts = findValidParts workflowMap (MetaPart sRange sRange sRange sRange) "in"
    sRange = Range 1 4000
    result = findCombinations validParts

type ID = String

data Range = Range { start :: Integer, end :: Integer } deriving (Show, Eq)

data MetaPart = MetaPart { x :: Range, m :: Range, a :: Range, s :: Range } deriving (Show, Eq)

data Workflow = Workflow { name :: ID, rules :: [Rule], final :: ID }

type Rule = (Char, Range, ID)

parseInput :: String -> [Workflow]
parseInput = map parseWorkflow . lines .  fst . splitOn "\n\n"

parseWorkflow :: String -> Workflow
parseWorkflow str = Workflow name rules final where
    name = takeWhile (/= '{') str
    innerStr = takeWhile (/='}') (drop (length name + 1) str)
    pieces = splitBy "," innerStr
    rules = map parseRule (take (length pieces - 1) pieces)
    final = pieces !! (length pieces - 1)

parseRule :: String -> Rule
parseRule str = (p, r, n) where
    (predStr, n) = splitOn ":" str
    p = head predStr
    v = read $ drop 2 predStr
    r = case predStr !! 1 of
        '<' -> Range 1 (v - 1)
        '>' -> Range (v + 1) 4000
        _ -> error "invalid rule"

findValidParts :: Map ID Workflow -> MetaPart -> ID -> [MetaPart]
findValidParts _ part "A" = [part]
findValidParts _ _ "R" = []
findValidParts workflowMap part name
    | isUselessPart part = []
    | otherwise = helper (rules workflow) part where
        workflow = workflowMap ! name
        helper :: [Rule] -> MetaPart -> [MetaPart]
        helper [] part = findValidParts workflowMap part (final workflow)
        helper (rule:xs) part = findValidParts workflowMap validPart validID ++ helper xs invalidPart where
            (validPart, invalidPart, validID) = splitPart rule part

splitPart :: Rule -> MetaPart -> (MetaPart, MetaPart, ID)
splitPart (component, range, validID) part = (validPart, invalidPart, validID) where
    validPart = case component of 
        'x' -> part { x = overlap range (x part) }
        'm' -> part { m = overlap range (m part) }
        'a' -> part { a = overlap range (a part) }
        's' -> part { s = overlap range (s part) }
    invalidPart = case component of 
        'x' -> part { x = overlap (inverseRule range) (x part) }
        'm' -> part { m = overlap (inverseRule range) (m part) }
        'a' -> part { a = overlap (inverseRule range) (a part) }
        's' -> part { s = overlap (inverseRule range) (s part) }

inverseRule :: Range -> Range
inverseRule (Range a b)
    | a == 1    = Range (b+1) 4000
    | b == 4000 = Range 1     (a-1)
    | otherwise = error "invalid range"

overlap :: Range -> Range -> Range
overlap (Range a b) (Range c d) = Range (max a c) (min b d)

isUselessPart :: MetaPart -> Bool
isUselessPart part = any (\(Range a b) -> b < a) [x part, m part, a part, s part]

findCombinations :: [MetaPart] -> Integer
findCombinations parts = result where
    result = sum $ map findComb parts
    findComb part = xSize * mSize * aSize * sSize where
        xSize = rangeSize (x part)
        mSize = rangeSize (m part)
        aSize = rangeSize (a part)
        sSize = rangeSize (s part)
    rangeSize (Range st en) = en - st + 1

I had a nice set of Workflow, Rule, and Part data structures for part 1, so I was able to add a PartRange for part 2 and extend the operations from Part 1 over the ranges.

The bug I hit for Part 2 is that I was returning a map of [outcome]PartRange from each workflow (where outcome is A, R or another workflow), and I didn't take into account the fact that there might be multiple results for the same outcome. Once I converted it to a list of [outcome, PartRange] pairs, everything else ran like clockwork.