AoC 2021 day 5

aoc2021d5.py

#!/usr/bin/env python3
#
# https://adventofcode.com/2021/day/5

def walk(start,end):
    if (start > end):
        return range(start,end-1,-1)
    return range(start,end+1)

class Segment:
    def __init__(self, line, diagonals=False):
        startstr,endstr = line.rstrip().split(" -> ")
        x1, y1 = map(int,startstr.split(","))
        x2, y2 = map(int,endstr.split(","))
        if x1 == x2:
            self.coords = { (x1, y) for y in walk(y1,y2) }
        elif y1 == y2:
            self.coords = { (x, y1) for x in walk(x1,x2) }
        elif diagonals:
            self.coords = set(zip(walk(x1,x2),walk(y1,y2)))
        else:
            self.coords = set() # diagonal but not counting

    def __contains__(self, other):
        if type(other) is Segment:
            return bool(self.coords & other.coords)
        else:
            return other in self.coords

assert Segment("0,0 -> 3,0").coords == { (0,0), (1,0), (2,0), (3,0) }
assert Segment("0,0 -> 0,3").coords == { (0,0), (0,1), (0,2), (0,3) }
assert Segment("0,0 -> 3,3").coords == set()
assert Segment("0,3 -> 3,0").coords == set()
assert Segment("0,0 -> 3,3",True).coords == { (0,0), (1,1), (2,2), (3,3) }
assert Segment("0,3 -> 3,0",True).coords == { (0,3), (1,2), (2,1), (3,0) }
assert (2,2) in Segment("1,2 -> 3,2")
assert (2,2) in Segment("2,1 -> 2,3")
assert (0,2) not in Segment("1,2 -> 3,2")
assert (2,0) not in Segment("1,2 -> 3,2")
assert Segment("1,2 -> 3,2") in Segment("2,1 -> 2,3")
assert Segment("1,2 -> 3,2") in Segment("0,2 -> 2,2")
assert Segment("1,2 -> 3,2") in Segment("0,2 -> 2,2")

def run(filename, diagonals=False):
    with open(filename) as fi:
        segments = {Segment(line,diagonals) for line in fi.readlines()}
    found = set()
    for s1 in segments:
        if not s1: continue
        for s2 in segments:
            if s1 is not s2 and s1 in s2:
                found.update(s1.coords & s2.coords)
    return len(found)

def show(filename, diagonals=False):
    with open(filename) as fi:
        segments = {Segment(line,diagonals) for line in fi.readlines()}
    for y in range(10):
        for x in range(10):
            c = len(list(filter( lambda z: (x,y) in z, segments )))
            print( ".123456"[c], end="")
        print("")

show("d5test.txt")
assert run("d5test.txt") == 5
print("part1:", run("d5data.txt"))
show("d5test.txt", True)
assert run("d5test.txt", True) == 12
print("part2:", run("d5data.txt", True))