ワーシャル-フロイド法のお勉強。最短路線検索。

warshalfloyd.py

# coding: utf8
from __future__ import division, print_function, unicode_literals
from functools import total_ordering
from collections import namedtuple
import copy


@total_ordering
class Way(namedtuple('Way', ('time', 'transporter'))):
    '''
    わざわざこんなクラスを用意してるけど、
    「大小比較さえできれば良い」ので、別に int でもよい。
    '''
    def __eq__(self, other):
        return sum(self.time) == sum(other.time)

    def __lt__(self, other):
        return sum(self.time) < sum(other.time)

    def __add__(self, other):
        return self._replace(
            time=self.time + other.time,
            transporter=self.transporter + other.transporter)

    @classmethod
    def atom(cls, time, transporter):
        return cls((time,), (transporter,))

    def format(self):
        return '計{} 時間: {}'.format(
            sum(self.time), '→'.join(self.transporter))


NO_WAY = Way.atom(float('inf'), 'たどりつかない')
NO_MOVE = Way.atom(0, 'そのまま')


def fill_adjacency(g):
    vertexes = g.keys()
    return {
        v1: {
            v2: NO_MOVE if v1 == v2 else g[v1].get(v2, NO_WAY)
            for v2 in vertexes
        } for v1 in vertexes
    }


def warshal_floyd(g):
    vertexes = g.keys()
    d = copy.deepcopy(g)
    for i in vertexes:
        for j in vertexes:
            for k in vertexes:
                d[j][k] = min(d[j][k], d[j][i] + d[i][k])
    return d


if __name__ == '__main__':
    '''
    言うまでもなく、この路線情報は架空のものですよ！
    マジレスだめよ!!! ネタレスならいいよ!!!
    '''
    gen = {
        '東京': {
            '仙台': Way.atom(4, '東北新幹線'),
            '新潟': Way.atom(2, '上越新幹線'),
            '名古屋': Way.atom(2, '東海道新幹線'),
        },
        '仙台': {'東京': Way.atom(4, '東北新幹線')},
        '新潟': {'東京': Way.atom(2, '上越新幹線')},
        '名古屋': {
            '東京': Way.atom(2, '東海道新幹線'),
            '大阪': Way.atom(1, 'リニア'),
        },
        '大阪': {'名古屋': Way.atom(1, 'リニア')},
        '火星': {},
    }
    t = warshal_floyd(fill_adjacency(gen))
    print(t['大阪']['仙台'].format())
    print(t['火星']['仙台'].format())