mines problem

mines.py

import re
import sys
import numpy as np


def read(stream=sys.stdin):
    '''Read some input to get a matrix of mines.

    Each row of the mine matrix is a different mine. The first column is the
    current gold value. The second column is the decay rate per travel step.

    Using a matrix means we can use matrix operations rather than loops.
    '''
    pat = re.compile('([0-9]+)\s+([0-9]+)')
    line = stream.readline()
    match = pat.match(line)
    n_mines = int(match[1])
    travel_time = int(match[2])
    mines = np.zeros((n_mines, 2))
    for i, line in enumerate(stream):
        match = pat.match(line)
        g = int(match[1])
        d = int(match[2])
        mines[i] = (g, d * travel_time)
    return mines


def generate(n):
    '''Generate a sample problem of size n.
    '''
    mines = np.ndarray((n, 2), dtype='int')
    for m in mines:
        m[0] = np.random.randint(2, 100)
        m[1] = np.random.randint(1, m[0])
    return mines


def advance(mines):
    '''Build a new mine matrix one step in the future.
    '''
    mines = np.copy(mines)
    mines[:,0] -= mines[:,1]
    mines[:,0].clip(0)
    return mines


def alive(mines):
    '''Returns True if we can move to a new mine.
    '''
    n_active = (mines[:,0] > 0).sum()
    return n_active > 1


def best(mines, exclude=[]):
    '''Returns the index of the best mine, excluding some list of indices.

    The best mine is the one with the most gold.
    Ties are broken by decay rate.
    '''
    m = np.copy(mines)
    m[exclude, 0] = 0
    gold = m[:, 0]
    decay = m[:, 1]
    m = gold * np.max(decay) + decay  # combine the main criteria and the tie breaker
    i = np.argmax(m)
    return i


def solve(mines, start=None):
    '''Returns the optimal score you can achieve on some mines problem.

    If no starting mine is given, start from the best mine.
    '''
    if start is None: start = best(mines)
    current = start
    score = 0
    while alive(mines):
        score += mines[current, 0]
        i = best(mines, exclude=current)
        assert current != i
        current = i
        mines = advance(mines)
    return score


if __name__ == '__main__':
    mines = read()
    x = solve(mines)
    print(x)