
# Python imports
from pathlib import Path
from typing import Callable, Dict, Iterable, List, Tuple

import numpy as np

# Paths to data
testpath = Path("day07_test.txt")
datapath = Path("day07_data.txt")


def load_input(fpath: Path) -> np.array:
    """Return list of crab horizontal positions
    
    :param fpath:  Path to data file
    """
    with fpath.open("r") as ifh:
        return np.array([int(_) for _ in ifh.read().strip().split(",")])


def simplecost(crabpos: np.array, testpos: int) -> int:
    """Returns cost function for simple crab movements
    
    :param crabpos:  array of initial crab positions
    :param testpos:  position for cost calculation    
    """
    return sum(abs(crabpos - testpos))


def optpos(crabpos: np.array, func: Callable) -> Tuple[int, int]:
    """Returns best position and fuel cost
    
    :param crabpos:  array of initial crab positions
    """
    
    # Initialise starting position (central), cost at
    # that position, and the stepsize (half the distance
    # from best position to crab position bounds)
    best_pos = int(0.5 * (crabpos.max() + crabpos.min()))
    mincost = func(crabpos, best_pos)
    stepsize = max([crabpos.max() - best_pos, best_pos - crabpos.min()])
    
    # At each iteration we establish a test position at
    # the point stepsize away above and below the current
    # best position.
    # If the cost at either test position is lower than
    # that at the best position, we move the best position
    # to that position.
    # If not, we reduce the stepsize by half.
    # We will halt when the stepsize falls below 1 - at
    # that point we have the minimum value
    while stepsize:
        # Establish test positions
        testpos_up, testpos_dn = (min(crabpos.max(), best_pos + stepsize),
                                  max(0, best_pos - stepsize))
        # Cost at test positions
        cost_up, cost_dn = (func(crabpos, testpos_up),
                            func(crabpos, testpos_dn))

        
        # We move if the cost at either position is lower
        if cost_up < mincost:
            best_pos = testpos_up
            mincost = cost_up
        elif cost_dn < mincost:
            best_pos = testpos_dn
            mincost = cost_dn
        else:  # else we halve the stepsize
            stepsize = int(0.5 * stepsize)

    # Return the optimal position and the cost at that position
    return best_pos, mincost


crabpos = load_input(testpath)
optpos(crabpos, simplecost)

(2, 37)


crabpos = load_input(datapath)
optpos(crabpos, simplecost)

(342, 351901)


# global cache of sums of 1 + 2 + ... + p_i,
# keyed by p_i
CACHE = {}

def movecost(dist: int) -> int:
    """Returns sum of 1+2+...+dist
    
    :param dist:  distance a crab has to travel
    
    Uses the global cache to get a value for
    dist, or calculates and updates the cache if
    not present.
    """
    if dist not in CACHE:
        # sum of first n integers is n * (n+1) / 2
        CACHE[dist] = int(0.5 * dist * (dist + 1))
    
    return CACHE[dist]

# Vectorise the movecost() function
movecost_v = np.vectorize(movecost)

def diffcost(crabpos: np.array, testpos: int) -> int:
    """Returns cost of crab movement
    
    :param crabpos:  array of initial crab positions
    :param testpos:  position for cost calculation
    
    Cost of movement is sum of 1+2+...+p_i for each p_i
    in the crab positions.
    """
    return sum(movecost_v(abs(crabpos - testpos)))


crabpos = load_input(testpath)
optpos(crabpos, diffcost)

(5, 168)


crabpos = load_input(datapath)
optpos(crabpos, diffcost)

(478, 101079875)

Day 07: The Treachery of Whales¶

Puzzle 2¶