ABC070: D - Transit Tree Path

問題

問題文

https://beta.atcoder.jp/contests/abc070/tasks/abc070_d

問題概要

頂点数が N の重み付き無向グラフが与えられる。

与えられるグラフは辺の数が N−1 本で閉路のない連結グラフ, すなわち木構造となっている。

頂点 から 頂点 K を経由しつつ, 頂点 まで移動する場合の最短経路を求めよ。

例

次のような入力が与えられたとき,

5
1 2 1
1 3 1
2 4 1
3 5 1
3 1
2 4
2 3
4 5

グラフは次のようになり,

求めるものは頂点 1 を経由して 2 から 4, 2 から 3, 4 から 5 へ移動する最短経路であるから,

3
2
4

が答えとなる。

解答例

指針

• 頂点 K から他の全頂点への最短経路を求める。

解説

グラフの表現

頂点数の制約 から隣接行列でグラフを持つと空間計算量が O(N²) となり MLE (Memory Limit Exceeded) することが分かる。

実際 10¹⁰ * 8 B = 80000 MB

でありこれはメモリ制限 256 MB を大幅に超える。

隣接リストを使うと空間計算量が O(N) となり制限内に収まる。

実装は std::vecotor を使うと楽である。

最短経路の求め方

頂点 K を経由したときの2点間の最短経路は頂点 K から他の全頂点への最短経路を求めればいいことが分かる。このような問題を単一始点最短路問題という。

単一始点最短路問題を解くアルゴリズムとして Bellman-Ford 法や Dijkstra 法が知られている。

• Bellman-Ford 法

  • 計算量 O(|V||E|)

• Dijkstra 法

  • 計算量 O (|E|log|V|)
  • 辺のコストが非負のときしか使えない

※ Dijkstra の計算量は実装方法によって異なる。

|V| は頂点数, |E| は辺の数を表している。

制約を考えると Bellman-Ford では無理で Dijkstra なら行けそうである。

しかし今回与えられたグラフは閉路がない連結グラフであるため, 任意の 2 頂点間の最短路は 1 通りに定まるので深さ優先探索 (Depth First Search: DFS) あるいは幅優先探索 (Breadth First Search: BFS) を1度行うだけで最短路が求まる。

DFS および BFS の計算量は O (|V| + |E|) で Dijkstra よりも高速である。

DFS による解法

• C++

#include <iostream>
#include <vector>
#define MAX_V 100001

using namespace std;

struct edge { int to; long long cost; };

vector<edge> tree[MAX_V];
long long depth[MAX_V];

// v: 今いる頂点, p: v の親
void dfs(int v, int p, long long d) {
    depth[v] = d;
    for (int i = 0; i < tree[v].size(); i++) {
        edge e = tree[v][i];
        // 戻らないようにする処理
        if (e.to == p) { continue; }
        dfs(e.to, v, d + e.cost);
    }
}

int main() {
    int N, a, b, c;
    int Q, K, x, y;

    cin >> N;

    for (int i = 0; i < N-1; i++) {
        cin >> a >> b >> c;
        a--, b--;
        tree[a].push_back(edge{b, c});
        tree[b].push_back(edge{a, c});
    }

    cin >> Q >> K;
    dfs(K-1, -1, 0);

    for (int i = 0; i < Q; i++) {
        cin >> x >> y;
        x--; y--;
        cout << depth[x] + depth[y] << endl;
    }
    return 0;
}

• Python

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import sys

sys.setrecursionlimit(200000)

N = int(raw_input())
tree = [[] for _ in xrange(N)]

for i in range(N-1):
    a, b, c = map(int, raw_input().split())
    a -= 1; b -= 1;
    tree[a].append((b, c))
    tree[b].append((a, c))

Q, K = map(int, raw_input().split())
K -= 1

depth = [0 for _ in xrange(N)]

def dfs(v, p = -1, d = 0):
    depth[v] = d
    for i, j in tree[v]:
        if i == p:
            continue
        dfs(i, v, d + j)

dfs(K)

for i in range(Q):
    x, y = map(int, raw_input().split())
    x -= 1; y -= 1
    print depth[x] + depth[y]

スタックを明示的に使った DFS による解法

• C++

#include <iostream>
#include <vector>
#include <stack>

#define MAX_V 100001

using namespace std;

struct edge { int to; long long cost; };
struct sets { int v, p; long long d; };
vector<edge> tree[MAX_V];
long long depth[MAX_V];

void dfs(int v, int p, long long d) {
    stack<sets> St;
    St.push(sets{v, p, d});

    while (!St.empty()) {
        sets cur = St.top(); St.pop();
        depth[cur.v] = cur.d;
        for (int i = 0; i < tree[cur.v].size(); i++) {
            edge e = tree[cur.v][i];
            if (e.to == cur.p) { continue; }
            St.push(sets{e.to, cur.v, cur.d + e.cost});
        }
    }
}

int main() {
    int N, a, b, c;
    int Q, K, x, y;

    cin >> N;

    for (int i = 0; i < N-1; i++) {
        cin >> a >> b >> c;
        a--, b--;
        tree[a].push_back(edge{b, c});
        tree[b].push_back(edge{a, c});
    }

    cin >> Q >> K; K--;
    dfs(K, -1, 0);

    for (int i = 0; i < Q; i++) {
        cin >> x >> y;
        x--; y--;
        cout << depth[x] + depth[y] << endl;
    }
    return 0;
}

• Python

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

N = int(raw_input())
tree = [[] for _ in xrange(N)]

for i in range(N-1):
    a, b, c = map(int, raw_input().split())
    a -= 1; b -= 1;
    tree[a].append((b, c))
    tree[b].append((a, c))

Q, K = map(int, raw_input().split())
K -= 1

depth = [0 for _ in xrange(N)]

def dfs(v, p = -1, d = 0):
    St = []
    St.append([v, p, d])
    # 空のシーケンスは偽になる
    while St:
        cur = St.pop()
        depth[cur[0]] = cur[2]
        for i, j in tree[cur[0]]:
            if i == cur[1]:
                continue
            St.append([i, cur[0], cur[2] + j])

dfs(K)

for i in range(Q):
    x, y = map(int, raw_input().split())
    x -= 1; y -= 1
    print depth[x] + depth[y]

BFS による解法

Stack を明示的に使った DFS と BFS のコードは非常によく似ている。

• C++

#include <iostream>
#include <vector>
#include <queue>

#define MAX_V 100001

using namespace std;

struct edge { int to; long long cost; };
struct sets { int v, p; long long d; };
vector<edge> tree[MAX_V];
long long depth[MAX_V];

void bfs(int v, int p, long long d) {
    queue<sets> q;
    q.push(sets{v, p, d});

    while (!q.empty()) {
        sets cur = q.front(); q.pop();
        depth[cur.v] = cur.d;
        for (int i = 0; i < tree[cur.v].size(); i++) {
            edge e = tree[cur.v][i];
            if (e.to == cur.p) { continue; }
            q.push(sets{e.to, cur.v, cur.d + e.cost});
        }
    }
}

int main() {
    int N, a, b, c;
    int Q, K, x, y;

    cin >> N;

    for (int i = 0; i < N-1; i++) {
        cin >> a >> b >> c;
        a--; b--;
        tree[a].push_back(edge{b, c});
        tree[b].push_back(edge{a, c});
    }

    cin >> Q >> K; K--;
    bfs(K, -1, 0);

    for (int i = 0; i < Q; i++) {
        cin >> x >> y;
        x--; y--;
        cout << depth[x] + depth[y] << endl;
    }
    return 0;
}

• Python

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import Queue

N = int(raw_input())
tree = [[] for _ in xrange(N)]

for i in range(N-1):
    a, b, c = map(int, raw_input().split())
    a -= 1; b -= 1;
    tree[a].append((b, c))
    tree[b].append((a, c))

Q, K = map(int, raw_input().split())
K -= 1

depth = [0 for _ in xrange(N)]

def bfs(v, p = -1, d = 0):
    q = Queue.Queue()
    q.put([v, p, d])
    while not q.empty():
        cur = q.get()
        depth[cur[0]] = cur[2]
        for i, j in tree[cur[0]]:
            if i == cur[1]:
                continue
            q.put([i, cur[0], cur[2] + j])

bfs(K)

for i in range(Q):
    x, y = map(int, raw_input().split())
    x -= 1; y -= 1;
    print depth[x] + depth[y]

Dijkstra による解法

• C++

#include <iostream>
#include <vector>
#include <queue>

#define MAX_V 100001
#define INF 1e18

using namespace std;

struct edge { int to; long long cost; };
typedef pair<long long, int> P;  // first は最短距離, second は頂点番号

vector<edge> tree[MAX_V];
long long d[MAX_V];
int N;

// s: 開始地点の頂点
void dijkstra(int s) {
    // greater<P> を指定することで first が小さい順に取り出せるようにする
    priority_queue<P, vector<P>, greater<P> > que;
    fill(d, d + N, INF);
    d[s] = 0;
    que.push(P(0, s));

    while (!que.empty()) {
        P p = que.top(); que.pop();
        int v = p.second;
        if (d[v] < p.first) { continue; }
        for (int i = 0; i < tree[v].size(); i++) {
            edge e = tree[v][i];
            if (d[e.to] > d[v] + e.cost) {
                d[e.to] = d[v] + e.cost;
                que.push(P(d[e.to], e.to));
            }
        }
    }
}

int main() {
    int a, b, c;
    int Q, K, x, y;

    cin >> N;

    for (int i = 0; i < N-1; i++) {
        cin >> a >> b >> c;
        a--, b--;
        tree[a].push_back(edge{b, c});
        tree[b].push_back(edge{a, c});
    }

    cin >> Q >> K; K--;

    dijkstra(K);

    for (int i = 0; i < Q; i++) {
        cin >> x >> y;
        x--; y--;
        cout << d[x] + d[y] << endl;
    }
    return 0;
}

• Python

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import Queue

N = int(raw_input())
tree = [[] for _ in xrange(N)]

for i in range(N-1):
    a, b, c = map(int, raw_input().split())
    a -= 1; b -= 1;
    tree[a].append((b, c))
    tree[b].append((a, c))

Q, K = map(int, raw_input().split())
K -= 1

d = [float("inf") for _ in xrange(N)]

def dijkstra(s):
    que = Queue.PriorityQueue()
    d[s] = 0
    que.put([0, s])

    while not que.empty():
        p = que.get()
        v = p[1]
        if d[v] < p[0]:
            continue
        for i, j in tree[v]:
            if d[i] > d[v] + j:
                d[i] = d[v] + j
                que.put([d[i], i])

dijkstra(K)

for i in range(Q):
    x, y = map(int, raw_input().split())
    x -= 1; y -= 1
    print d[x] + d[y]

参考文献

• https://img.atcoder.jp/abc070/editorial.pdf

• 蟻本