CodeForces Round #303 Div2 E - Paths and Trees

問題 : http://codeforces.com/contest/545/problem/E

内容:

頂点 : , 辺 : から成る多重辺の無い無向グラフG(V,E)が与えられる．このとき，あるノードからの最短パスのみで構成される， コストの総和が最小の木を求めよ．

解法:

からの最短距離をとりあえず求めて，あるノードとその一つ前のノードとの距離が小さいエッジを集めていく方法で解いた． エッジそれぞれについてそのエッジを用いた場合の距離が最短距離であり且つ一つ前の頂点とのコストが小さいものをとっていく．

具体的には， dist[u] == dist[v] + w and w < dist[u] - dist[prev[u]] となるようなエッジを採用していく． 全体でとなりくらいの計算になるので，PyPyでは間に合わないみたい．

雑感:

C++に書き直して通りはしたけど，ソースは汚い上に筋が悪そう． 今回，コストが最大であることを忘れてintのまま計算を行っており相当WAを出してしまった．C++のソースの見通しが悪すぎるのでTLEになったPythonのコードも貼っておく．読みやすいコードが書けるようになりたい．

TLEになったPythonコード :

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

from heapq import heappush, heappop
def dijkstra(adjList,s):
    num = len(adjList)
    dist = [10**10 for i in range(num)]
    prev = [-1 for i in range(num)]
    queue = []

    heappush(queue,(0,s))
    dist[s] = 0
    prev[s] = (0,-1)
    while queue != []:
        v_cost, v = heappop(queue)
        if dist[v] < v_cost:
            continue
        for u_cost, u, edge_no in adjList[v]:
            if u != v and u_cost + dist[v] < dist[u]:
                prev[u] = (v,edge_no)
                dist[u] = u_cost + dist[v]
                heappush(queue, (dist[u], u))
    return dist,prev

n,m = map(int,input().split())
adjList = [[] for i in range(n)]
edges = []
for i in range(m):
    u,v,w = map(int,input().split())
    u-=1
    v-=1
    adjList[u].append((w,v,i))
    adjList[v].append((w,u,i))
    edges.append((u,v,w))

s = int(input())-1
dist,prev = dijkstra(adjList,s)
idx = 0
for u,v,w in edges:
    if dist[u] == dist[v] + w and w < dist[u] - dist[prev[u][0]]:
        prev[u] = (v,idx)
    elif dist[v] == dist[u] + w and w < dist[v] - dist[prev[v][0]]:
        prev[v] = (u,idx)
    idx += 1

ans = 0
ans_edges = []
for node,idx in prev:
    if idx >= 0:
        ans += edges[idx][2]
        ans_edges.append(idx+1)
print(ans)
print(' '.join(map(str,sorted(ans_edges))))

通ったc++のコード:

#include <bits/stdc++.h>

using namespace std;

typedef vector<int> vi;
typedef pair<int, int> pii;
typedef long long ll;

#define FOR(i,a,b) for(int i=(a);i<(b);i++)
#define REP(i,n)  FOR(i,0,n)
#define INF 1e15

struct edge{
    int to;
    ll cost;

    bool operator>(const edge &e)const{
        return cost > e.cost;
    }

    bool operator<(const edge &e)const{
        return cost < e.cost;
    }
};

void dijkstra(vector<vector<pair<edge,int>> > &g, vector<ll> &dist, vector<pii> &prev, int s){
    priority_queue<edge, vector<edge>, greater<edge> > pq;

    fill(dist.begin(),dist.end(),INF);
    dist[s] = 0;
    prev[s] = (pii){0,-1};

    pq.push((edge){s,0});

    while(!pq.empty()){
        auto v = pq.top(); pq.pop();
        if(dist[v.to] < v.cost) continue;

        for(pair<edge,int> u:g[v.to]){
            edge next_v = u.first;
            if(next_v.to != v.to and next_v.cost + dist[v.to] < dist[next_v.to]){
                prev[next_v.to] = pii{v.to,u.second};
                dist[next_v.to] = next_v.cost + dist[v.to];
                pq.push((edge){next_v.to,dist[next_v.to]});
            }
        }
    }
}

int main(){
    int n,m;
    cin >> n >> m ;
    vector<vector<pair<edge,int>>> G(n);
    typedef tuple<int,int,int> tiii;
    vector<tiii> edges;
    FOR(i,0,m){
        int u,v,w;
        cin >> u >> v >> w ;
        u--;v--;
        G[u].push_back((pair<edge,int>{(edge){v,w},i}));
        G[v].push_back((pair<edge,int>((edge){u,w},i)));
        edges.push_back((tiii){u,v,w});
    }

    int s;
    cin >> s ;

    vector<ll> dist(n);
    vector<pii> prev(n);
    dijkstra(G,dist,prev,s-1);

    FOR(i,0,m){
        int u,v,w;
        tie(u,v,w) = edges[i];
        if(dist[v] == (dist[u] + w) && w < (dist[v] - dist[prev[v].first]))
            prev[v] = pii{u,i};
        if(dist[u] == (dist[v] + w) && w < (dist[u] - dist[prev[u].first]))
            prev[u] = pii{v,i};
    }

    ll ans = 0;
    vector<int> ans_edge;
    for(auto p : prev){
        int node,idx;
        tie(node,idx) = p;
        if(idx >= 0){
            int u,v,w;
            tie(u,v,w) = edges[idx];
            ans += w;
            ans_edge.push_back(idx+1);
        }
    }

    cout << ans << endl;
    sort(ans_edge.begin(),ans_edge.end());
    for(auto i:ans_edge)
        cout << i << " ";
    cout << endl;

    return 0;
}