Solution to Advent of Code 2016 Day 1

01-city-block-distance.py

import numpy as np
import pytest as t
import sys

TURN = {
    'L': np.array([[0,  1],
                   [-1, 0]]),
    'R': np.array([[0, -1],
                   [1,  0]])
}
ORIGIN = np.array([0, 0])
NORTH = np.array([0, 1])


class Santa:
    def __init__(self, location, heading):
        self.location = np.array(location)
        self.heading = np.array(heading)

    def execute_one(self, instruction):
        self.heading = self.heading @ TURN[instruction[0]]
        self.location += self.heading * int(instruction[1:])

    def execute_many(self, instructions):
        for i in instructions.split(','):
            self.execute_one(i.strip())

    def distance_from_start(self):
        return sum(abs(self.location))

    def __str__(self):
        return f'Santa @ {self.location}, heading {self.heading}'


def test_execute_one():
    s = Santa(ORIGIN, NORTH)

    s.execute_one('L1')
    assert all(s.location == np.array([-1, 0]))
    assert all(s.heading == np.array([-1, 0]))

    s.execute_one('L3')
    assert all(s.location == np.array([-1, -3]))
    assert all(s.heading == np.array([0, -1]))

    s.execute_one('R3')
    assert all(s.location == np.array([-4, -3]))
    assert all(s.heading == np.array([-1, 0]))

    s.execute_one('R100')
    assert all(s.location == np.array([-4, 97]))
    assert all(s.heading == np.array([0, 1]))


def test_execute_many():
    s = Santa(ORIGIN, NORTH)

    s.execute_many('L1, L3, R3')
    assert all(s.location == np.array([-4, -3]))
    assert all(s.heading == np.array([-1, 0]))

def test_distance():
    assert Santa(ORIGIN, NORTH).distance_from_start() == 0
    assert Santa((10, 10), NORTH).distance_from_start() == 20
    assert Santa((-17, 10), NORTH).distance_from_start() == 27


def test_turn_left():
    east = NORTH @ TURN['L']
    south = east @ TURN['L']
    west = south @ TURN['L']

    assert all(east == np.array([-1, 0]))
    assert all(south == np.array([0, -1]))
    assert all(west == np.array([1, 0]))


def test_turn_right():
    west = NORTH @ TURN['R']
    south = west @ TURN['R']
    east = south @ TURN['R']

    assert all(east == np.array([-1, 0]))
    assert all(south == np.array([0, -1]))
    assert all(west == np.array([1, 0]))


if __name__ == '__main__':
    instructions = sys.stdin.read()

    santa = Santa(ORIGIN, NORTH)
    santa.execute_many(instructions)

    print(santa)
    print(santa.distance_from_start())