PythonでAizu Online Judgeを解く/GRL_1_A - Single Shortest Path

Posted on

概要

重み付きの有向グラフ \(G(V, E)\) が与えられます。頂点 \(r\) を始点としたときの各頂点までの最短経路の長さを出力してください。

問題URL:
https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/1/GRL_1_A

制約

  • \(1 \leq |V| \leq 100000\)
  • \(0 \leq d_i \leq 10000\)
  • \(0 \leq |E| \leq 500000\)
  • \(G\)に平行辺はない
  • \(G\)に自己ループはない

Input

%python
def input(f):
    global v, e, r, s, t, d, g
    
    v, e, r = map(int, f.readline().split())
    s = []
    t = []
    d = []
    for _ in range(e):
        s_, t_, d_ = map(int, f.readline().split())
        s.append(s_)
        t.append(t_)
        d.append(d_)
        
    g = [list() for _ in range(v)]
    for i in range(e):
        g[s[i]].append(i)

Example Input

%sh
cat << EOF > /tmp/input1
4 5 0
0 1 1
0 2 4
1 2 2
2 3 1
1 3 5
EOF

cat << EOF > /tmp/input2
4 6 1
0 1 1
0 2 4
2 0 1
1 2 2
3 1 1
3 2 5
EOF

Python: networkxで可視化してみる

%sh
pip install networkx
Collecting networkx
  Downloading https://files.pythonhosted.org/packages/f3/f4/7e20ef40b118478191cec0b58c3192f822cace858c19505c7670961b76b2/networkx-2.2.zip (1.7MB)
Collecting decorator>=4.3.0 (from networkx)
  Downloading https://files.pythonhosted.org/packages/bc/bb/a24838832ba35baf52f32ab1a49b906b5f82fb7c76b2f6a7e35e140bac30/decorator-4.3.0-py2.py3-none-any.whl
Building wheels for collected packages: networkx
  Running setup.py bdist_wheel for networkx: started
  Running setup.py bdist_wheel for networkx: finished with status 'done'
  Stored in directory: /root/.cache/pip/wheels/68/f8/29/b53346a112a07d30a5a84d53f19aeadaa1a474897c0423af91
Successfully built networkx
Installing collected packages: decorator, networkx
Successfully installed decorator-4.3.0 networkx-2.2
You are using pip version 9.0.1, however version 18.1 is available.
You should consider upgrading via the 'pip install --upgrade pip' command.


%python
import os

if os.environ.get('ZEPPELIN_HOME'):
    import networkx as nx
    import matplotlib.pyplot as plt
    
    def visualize():
        G = nx.DiGraph()
        for i in range(v):
            G.add_node(i)
        for i in range(e):
            G.add_edge(s[i], t[i], t=d[i])
        
        plt.figure()
        plt.style.use('seaborn')
        plt.axis("off")
        pos = nx.spring_layout(G)
        nodes = nx.draw_networkx_nodes(G, pos, node_color='w', linewidths=2)
        nx.draw_networkx_labels(G, pos)
        nodes.set_edgecolor('black')
        nx.draw_networkx_edges(G, pos)
        nx.draw_networkx_edge_labels(G, pos)
        plt.show()

%python
if os.environ.get('ZEPPELIN_HOME'):
    with open('/tmp/input1') as f:
        print('Example #1')
        input(f)
        visualize()
Example #1


%python
if os.environ.get('ZEPPELIN_HOME'):
    with open('/tmp/input2') as f:
        print('Example #2')
        input(f)
        visualize()
Example #2

考え方

隣接する頂点について幅優先探索していきます。Pythonではcollectionsモジュールのdequeを使うと良さそうでした。

%python
import sys
import collections

def solve():
    inf = sys.maxint
    q = collections.deque()
    q.append(r)
    mincost = [inf]*v
    mincost[r] = 0
    while len(q) > 0:
        cur = q.popleft()
        for e in g[cur]:
            if mincost[t[e]] == inf or mincost[t[e]] > mincost[cur] + d[e]:
                mincost[t[e]] = mincost[cur] + d[e]
                q.append(t[e])
    
    for i in range(v):
        if mincost[i] == inf:
            print('INF')
        else:
            print(mincost[i])

Example Output #1

%python
if os.environ.get('ZEPPELIN_HOME'):
    os.environ['INPUT'] = '/tmp/input1'

with open(os.environ.get('INPUT') or '/dev/stdin') as f:
    input(f)
    solve()
0
1
3
4

Example Output #2

%python
if os.environ.get('ZEPPELIN_HOME'):
    os.environ['INPUT'] = '/tmp/input2'

    with open(os.environ.get('INPUT') or '/dev/stdin') as f:
        input(f)
        solve()
3
0
2
INF


sh