aoc2021-05.py

#!/usr/bin/env python3

import collections
import dataclasses
import re
import sys

PART = 2

PATTERN = re.compile("([0-9]+),([0-9]+) -> ([0-9]+),([0-9]+)")


def sign(x):
    if x != 0:
        return x / abs(x)
    else:
        return 0


@dataclasses.dataclass
class Line:
    x1: int
    y1: int
    x2: int
    y2: int

    def iter_points(self):
        dx = self.x2 - self.x1
        dy = self.y2 - self.y1

        if PART == 1:
            # vertical or horizontal allowed
            if dx != 0 and dy != 0:
                return
        elif PART == 2:
            # vertical, horizontal or diagonal allowed
            if dx != 0 and dy != 0 and abs(dx) != abs(dy):
                return
        else:
            assert False, "The puzzle only has two parts"

        x, y = self.x1, self.y1
        sx = sign(dx)
        sy = sign(dy)

        while True:
            yield x, y
            if (x, y) == (self.x2, self.y2):
                # reached the end of the line
                break
            x += sx
            y += sy


lines = [PATTERN.match(line).groups() for line in sys.stdin]
lines = [Line(*[int(e) for e in line]) for line in lines]

counter = collections.Counter()
for line in lines:
    counter.update(line.iter_points())

mc = counter.most_common()
print(sum(1 for e, c in mc if c > 1))