aoc-day13.py

#!/usr/bin/env python3
from __future__ import annotations

import dataclasses
import math
import os
import time
from functools import cmp_to_key
from typing import Callable

import aocd
import parse

AOC_DAY = 13
test_data = """\
[1,1,3,1,1]
[1,1,5,1,1]

[[1],[2,3,4]]
[[1],4]

[9]
[[8,7,6]]

[[4,4],4,4]
[[4,4],4,4,4]

[7,7,7,7]
[7,7,7]

[]
[3]

[[[]]]
[[]]

[1,[2,[3,[4,[5,6,7]]]],8,9]
[1,[2,[3,[4,[5,6,0]]]],8,9]
"""


def data_parse(dstr):
    return dstr.split("\n\n")


test_data = data_parse(test_data)

assert len(test_data) == 8

test_answer_one = 13
test_answer_two = 140


def split_eval(dstr) -> tuple[list, list]:
    left, right = dstr.strip().split("\n")
    return eval(left), eval(right)


def compare_list(left, right):
    for lv, rv in zip(left, right):
        ttype = (type(lv), type(rv))

        if ttype == (int, int):
            v = compare_int(lv, rv)
        elif ttype == (list, list):
            v = compare_list(lv, rv)
        elif ttype == (list, int):
            v = compare_list(lv, [rv])
        elif ttype == (int, list):
            v = compare_list([lv], rv)
        else:
            raise
        if v is not None:
            return v
    # out of options.
    dlen = len(left) - len(right)

    if dlen == 0:
        return None
    else:
        return dlen < 0


def compare_int(left, right):
    if left == right:
        return None
    else:
        return left < right


def get_correct_pair_idx(data) -> list[int]:
    correct_pairs = []
    for idx, pair in enumerate(data):
        idx += 1  # index starts at one
        left_data, right_data = split_eval(pair)

        if compare_list(left_data, right_data) is True:
            correct_pairs.append(idx)

    return correct_pairs


def part_one(data: list[list[str]]) -> int:
    pairs = get_correct_pair_idx(data)
    print(pairs)
    return sum(pairs)


def part_two(data: list[str]) -> int:
    output_pairs = [[[2]], [[6]]]

    # add pairs from input
    for pair in data:
        left_data, right_data = split_eval(pair)
        output_pairs.append(left_data)
        output_pairs.append(right_data)

    def intify(lv, lr):
        v = compare_list(lv, lr)
        if v is None:
            return 0
        return 1 if v else -1

    # sort, using our compare list function:
    sorted_list = sorted(output_pairs, key=cmp_to_key(intify), reverse=True)
    for x in sorted_list:
        print(x)

    idx2 = sorted_list.index([[2]]) + 1
    idx6 = sorted_list.index([[6]]) + 1

    return idx2 * idx6


if __name__ == "__main__":

    part_one_ans = part_one(test_data)
    assert part_one_ans == test_answer_one, f"{part_one_ans=}, not {test_answer_one=}"

    real_data = data_parse(aocd.get_data(day=AOC_DAY, year=2022))
    print("part 1:", part_one(real_data))

    part_two_ans = part_two(test_data)
    assert part_two_ans == test_answer_two, f"{part_two_ans=}, not {test_answer_two=}"
    print("part 2:", part_two(real_data))