Learning Algorithms

アルゴリズムの勉強メモ

Atcoder ARC #048 D. たこ焼き屋とQ人の高橋君

Atcoder ARC #048 D. たこ焼き屋とQ人の高橋君

D: たこ焼き屋とQ人の高橋君 - AtCoder Regular Contest 048 | AtCoder

ある木と、各ノードにたこ焼き屋さんがあるかどうかの情報が与えられる。最初はノード間の移動に2秒かかるが、たこ焼き屋さんがあるノードでたこ焼きを食べれば、1秒でノード間を移動できるようになる。クエリ(s, t)が順に与えられるので、sからtまで移動するのにかかる最短の時間を求めていく、という問題。まず木なので、(s, t)の最短経路は必ず通る。たこ焼きを食べる場合は、その経路上のあるノードからたこ焼き屋さんまでの経路を通り、折り返し、もとのノードに戻り、もとの経路を進むことになるはずである。そのあるノードをx、たこ焼き屋さんをyとし、xとyの距離をd(x, y)で表すと、求める値は、min(2 * d(s, x) + 2 * d(x, y) + d(x, y) + d(x, t)) = min(d(s, t) + d(s, x) + 3 * d(x, y))である。下準備として、各ノードについて、たこ焼き屋さんまでの最短距離pを保存しておく。これは、たこ焼き屋さんのあるノードをキューにあらかじめp == 0 としながらpushしておき、そこからp == INFならば更新していく幅優先探索することで調べられる。また、様々な距離を扱うため、根付き木で考えると都合が良い。根付き木上で、ごにゃごにゃ距離を計算すると、3 * p[i] + depth[i]と3 * p[i] - depth[i]の最小値が重要になることがわかる。木の任意のパス上の頂点に書かれた数の最小値を高速で計算するために、ダブリングとdpをうまく使う(HL分解的なことをして、セグメント木にのせたりできるらしいが、それはよく知らない)。dp[i][j] := ノードiから、距離の2のj乗まで親をたどったときのその区間の最小値、をあらかじめ計算しておく。すると、log(n)で各パス上の最小値を計算できる。図なしには説明しにくいが、lca(s, t)の位置とxの位置の前後関係で場合分けをしてそれぞれ計算する。

#include "bits/stdc++.h"
using namespace std;
#define rep(i, n)             for (int (i) = 0; (i) < (int)(n); (i)++)
template<typename T, typename U> inline void amin(T &x, U y) { if (y < x) x = y; }
 
#define maxv 101010
#define maxl 20
static const int INF = 0x3f3f3f; //大きくしすぎるとオーバーフローするので注意
 
int depth[maxv];
int dbl[maxv][maxl];
int dp[2][maxv][maxl];
vector<int> g[maxv];
 
int lca(int x, int y) {
        if (depth[x] < depth[y]) swap(x, y);
        rep(i, maxl) { 
                if ((depth[x] - depth[y]) & (1 << i)) x = dbl[x][i];
        }
        if (x == y) return x;
        for (int i = maxl - 1; i >= 0; i --) {
                if (dbl[x][i] != dbl[y][i]) {
                        x = dbl[x][i];
                        y = dbl[y][i];
                }
        }
        return dbl[x][0];
}
 
int getmin(int x, int d, int t) { //木上において、ノードxから距離dだけ親をたどったところまでの区間の最小値をダブリングで求める
        int res = INF;
        rep(i, maxl) {
                if (d & (1 << i)) {
                        res = min(res, dp[t][x][i]);
                        x = dbl[x][i];
                }
        }
        return res;
}
 
void dfs(int v, int u, int d) {
        dbl[v][0] = u;
        depth[v] = d;
        for (auto i : g[v]) {
                if (u == i) continue;
                dfs(i, v, d + 1);
        }
}
 
signed main() { 
        int n, query;
        cin >> n >> query;
        rep(i, n - 1) {
                int a, b;
                cin >> a >> b;
                a --, b --;
                g[a].push_back(b);
                g[b].push_back(a);
        }
        string s;
        cin >> s;
        dfs(0, 0, 0);
        int p[maxv];
        rep(i, n) p[i] = INF; //最初はすべてINF
        queue<int> q;
        rep(i, n) if (s[i] == '1') {
                q.push(i);
                p[i] = 0;
        }
        while (!q.empty()) { //たこ焼き屋さんがあるところから幅優先探索
                int x = q.front(); q.pop();
                for (auto i : g[x]) if (p[i] == INF) {
                        q.push(i);
                        p[i] = p[x] + 1;
                }
        }
        rep(i, maxl - 1) rep(j, n) dbl[j][i + 1] = dbl[dbl[j][i]][i]; //ダブリングの初期化
        rep(i, n) rep(t, 2) dp[t][i][0] = p[i] * 3 + (t == 0 ? (-1) * depth[i] : depth[i]);
        rep(i, maxl - 1) rep(j, n) rep(t, 2) dp[t][j][i + 1] = min(dp[t][j][i], dp[t][dbl[j][i]][i]);
        while (query --) {
                int s, t;
                cin >> s >> t;
                s --, t --;
                int c = lca(s, t);
                int l = depth[s] - depth[c], r = depth[t] - depth[c];
                int x = getmin(s, l, 0), y = getmin(t, r, 1);
                int ans = (l + r) * 2; //寄り道しない場合
                amin(ans, l * 2 + p[c] * 3 + r); //lcaがxと一致するとき
                amin(ans, l * 2 + x + depth[c] + r); //lcaよりs側にxがあるとき
                amin(ans, l * 2 + y + r - depth[c]); //lcaよりt側にxがあるとき
                cout << ans << endl;
        }
        return 0;
}               

dpが難しい。。。。