ITエンジニアのブログ

IT企業でエンジニアやってる人間の日常について

ARC015 [C]変わった単位

競技プログラミングに慣れていないので、幅優先探索深さ優先探索を素早く記述することができないし、ダイクストラについては擬似アルゴリズムを見ないとかけないという状態なので練習します。
AtCoder Regular Contest 15 の C 問題「変わった単位」
一行が次の内容で構成される複数行の入力に対し、一番大きい単位を一番小さい単位で表記しろという問題。
large ratio small
(例: week 7 day)
解説を見てみると、幅優先探索やワーシャルフロイトを使って一つの単位を全ての単位で表せば、単位の大小と比率が分かるとあった。
何となく幅優先探索ダイクストラで出来そうだと思うも、数年書いてないダイクストラアルゴリズムを忘れてしまったので Wikipedia を参照しました。
まずは単純なダイクストラのプログラム、

def gets():
    try:
        return input()
    except EOFError:
        return None

def main():
    Infinity = 1000000000
    start = input().strip()
    stations = []
    rails = []
    dists = {}
    que = []

    while True:
        line = gets()

        if not line:
            break

        xs = line.split()
        s = xs[0]
        c = int(xs[1])
        g = xs[2]

        if s not in stations:
            stations.append(s)
        if g not in stations:
            stations.append(g)

        rails.append((s, c, g))

    for s in stations:
        dists[s] = Infinity
        que.append(s)

    dists[start] = 0

    while len(que) > 0:
        minval = Infinity
        minkey = None

        for key in que:
            if minval > dists[key]:
                minval = dists[key]
                minkey = key

        que.remove(minkey)

        for (s, c, g) in [(s, c, g) for (s, c, g) in rails if s == minkey]:
            if dists[g] > dists[minkey] + c:
                dists[g] = dists[minkey] + c
        for (s, c, g) in [(s, c, g) for (s, c, g) in rails if g == minkey]:
            if dists[s] > dists[minkey] + c:
                dists[s] = dists[minkey] + c

    print(dists)

if __name__ == '__main__':
    main()

出発の駅と各駅間のコストを入力として受け取り、出発の駅から各駅までの最小コストを出力するプログラムです。例えばこのように入力を与えると、

Shibuya
Shibuya 3 Shinjuku
Ikebukuro 9 Shinjuku
Shibuya 4 Meguro
Shinagawa 6 Shibuya
Shinagawa 3 Shinbashi
Shinbashi 5 Tokyo
Tokyo 7 Ueno
Shinjuku 2 Ueno
Akabane 8 Ueno
Ueno 4 Shinagawa
Shinjuku 5 Meguro

このように出力します。

{'Ikebukuro': 12, 'Tokyo': 12, 'Shibuya': 0, 'Akabane': 13, 'Shinagawa': 6, 'Meguro': 4, 'Shinjuku': 3, 'Shinbashi': 9, 'Ueno': 5}

ARC015 の C 問題は到達可能問題に比率のデータをもたせれば良いことに気づいたので、ただの幅優先探索になりました。
単位を頂点、単位と比率と頂点の組を辺とする。
キューに一つの最初の頂点を入れる。

  1. 頂点から一つ頂点v0を取り出す。
  2. 辺の端にv0を持つものの反対の頂点v1とする。
  3. 比率が求まっていなければ頂点をキューに挿入し比率を求める。

を繰り返す。で、一つの単位を全ての単位で表現することが出来ます。
プログラムはこうなりました。

def main():
    n = int(input())
    units = []
    edges = []
    ratios = {}
    queue = []

    for i in range(n):
        xs = input().split()
        large = xs[0]
        small = xs[2]
        if large not in units:
            units.append(large)
        if small not in units:
            units.append(small)
        edges.append((large, int(xs[1]), small))

    start = units[0]

    for key in units:
        ratios[key] = None

    ratios[start] = 1

    queue.append(start)

    while len(queue) > 0:
        u = queue.pop(0)

        for (s, c, g) in [(s, c, g) for (s, c, g) in edges if s == u]:
            if ratios[g] == None:
                ratios[g] = ratios[u] / c
                queue.append(g)
        for (s, c, g) in [(s, c, g) for (s, c, g) in edges if g == u]:
            if ratios[s] == None:
                ratios[s] = ratios[u] * c
                queue.append(s)

    min_key = None
    min_value = 10e20
    max_key = None
    max_value = -10e20

    for key in ratios:
        if ratios[key] < min_value:
            min_value = ratios[key]
            min_key = key
        if ratios[key] > max_value:
            max_value = ratios[key]
            max_key = key

    print("{0}{1}={2}{3}".format(1, max_key, int(ratios[max_key] / ratios[min_key] + 0.5), min_key))

if __name__ == "__main__":
    main()

できれば C++ で書けるようになりたいのだが、 C++ の priority_queue をよく知らないので勉強しなければならない。というか STL の挙動全般についてあまり理解できていないのでもっと C++STL を勉強したい。