summaryrefslogtreecommitdiffstats
path: root/aoc/datastructures.py
blob: 244aa1e9744eda54f3e93a84d8a333c20eac74b5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
# -*- coding: utf-8 -*-
from collections import namedtuple
from typing import Iterator


class Coordinate(namedtuple("Coordinate", ["x", "y"])):
    def __sub__(self, other: "Coordinate") -> "Coordinate":
        return Coordinate(self.x - other.x, self.y - other.y)

    def __add__(self, other: "Coordinate") -> "Coordinate":
        return Coordinate(self.x + other.x, self.y + other.y)

    def manhattan_distance(self, other: "Coordinate") -> int:
        return abs(self.x - other.x) + abs(self.y - other.y)

    @property
    def polarity(self) -> "Coordinate":
        try:
            px = abs(self.x) / self.x
        except ZeroDivisionError:
            px = 0

        try:
            py = abs(self.y) / self.y
        except ZeroDivisionError:
            py = 0

        return Coordinate(
            px,
            py,
        )

    def neighbours(self, no_diagonal: bool = False) -> Iterator["Coordinate"]:
        if no_diagonal:
            yield self + Coordinate(-1, 0),
            yield self + Coordinate(1, 0),
            yield self + Coordinate(0, -1),
            yield self + Coordinate(0, 1),
        else:
            for dy in range(self.y - 1, self.y + 2):
                for dx in range(self.x - 1, self.x + 2):
                    if dy == self.y and dx == self.x:
                        continue

                    yield Coordinate(dx, dy)


class Coordinate3(namedtuple("Coordinate3", ["x", "y", "z"])):
    def __sub__(self, other: "Coordinate3") -> "Coordinate3":
        return Coordinate3(self.x - other.x, self.y - other.y, self.z - other.z)

    def __add__(self, other: "Coordinate3") -> "Coordinate3":
        return Coordinate3(self.x + other.x, self.y + other.y, self.z + other.z)

    def manhattan_distance(self, other: "Coordinate3") -> int:
        return abs(self.x - other.x) + abs(self.y - other.y) + abs(self.z - other.z)

    @property
    def polarity(self) -> "Coordinate3":
        try:
            px = abs(self.x) / self.x
        except ZeroDivisionError:
            px = 0

        try:
            py = abs(self.y) / self.y
        except ZeroDivisionError:
            py = 0

        try:
            pz = abs(self.z) / self.z
        except ZeroDivisionError:
            pz = 0

        return Coordinate3(
            px,
            py,
            pz,
        )

    def neighbours(self, no_diagonal: bool = False) -> Iterator["Coordinate3"]:
        if no_diagonal:
            yield self + Coordinate3(-1, 0, 0)
            yield self + Coordinate3(1, 0, 0)
            yield self + Coordinate3(0, -1, 0)
            yield self + Coordinate3(0, 1, 0)
            yield self + Coordinate3(0, 0, -1)
            yield self + Coordinate3(0, 0, 1)
        else:
            for dz in range(self.z - 1, self.z + 2):
                for dy in range(self.y - 1, self.y + 2):
                    for dx in range(self.x - 1, self.x + 2):
                        coordinate = Coordinate3(dx, dy, dz)

                        if dy == self.y and dx == self.x:
                            continue

                        yield coordinate