日々drdrする人のメモ

今日もdrdr、明日もdrdr

CodeChef July Challenge 2018: Subway Ride

CodeChef July Challenge 2018の問題: Subway Ride (Code: SUBWAY)
問題ページ: https://www.codechef.com/JULY18A/problems/SUBWAY

問題概要

頂点 {N}、辺 {M}の無向グラフが与えられる。
このグラフは自己ループと単純閉路が存在せず、多重辺を持つ。(つまり、同じ頂点同士を繋ぐ多重辺を全て1つの辺に置き換えると木になる)
そして、各辺には色が塗られている。

 {i}番目の辺は頂点 {u_i}と頂点 {v_i}を繋ぎ、 {w_i}の色で塗られている。


このグラフの2点間の単純パスを考え、この時通る辺の色を通った順に並べ、コストを計算する。
このコストは、単純パスで通った辺の内、隣接している辺の色が異なる数、で計算する。

今回、 {Q}クエリが与えられる。
各クエリで頂点 {u}と頂点 {v}が与えられるので、その2点間の単純パスで考えられる最大コストを計算せよ。

制約

 {1 \le N, Q \le 5 \times 10^5}
 {1 \le M \le 10^6}
 {1 \le u_i, v_i \le N}
 {u_i \neq v_i}
 {1 \le w_i \le M}

解法

LCAのようにダブリングしながらコストをdpで計算し、LCAを用いながら各クエリの答えを計算する。


まず、ある2点間に3色以上の多重辺が存在する場合、単純パス内にて、この辺と隣接する両端の辺の色と異なる色を必ず通ることができる
そのため、3色以上の多重辺は、どの色とも必ず異なる色を持つ1つの辺と見なすことができるため、多重辺を高々2つにすることができる。


あとは、与えられる無向グラフを木とみなし、LCAのようにダブリングしながら、最大となる時のコストをdpで計算していけばよい。

具体的には、

 {dp[v][k][c_1][c_2] = }各頂点 {v}から {2^k}上の親頂点 {w}の間の単純パスについて、 {v}から出る際に通った辺のうち色 {c_1}を通り、かつ色 {w}に到達する直前の辺のうち {c_2}を通った時の最大コスト

とする。下図のように頂点 {v, w}間両端の辺の色を決めて最大コストを求める感じ。
f:id:smijake3:20180719004638p:plain

各最大コストの計算は、 {w} {v}から {2^k}上の親頂点とすると、

 {\displaystyle dp[v][k+1][c_1][c_4] = \max_{c_2, c_3}(dp[v][k][c_1][c_2] + dp[w][k][c_3][c_4] + (c_2 \neq c_3))}

で行える。( {(c_2 \neq c_3)} {c_2} {c_3}の色が異なれば1, 同じであれば0とする)


そして、各クエリで与えられる2頂点 {u, v}間の最大コストは、ダブリングで2頂点のLCAを求めるように計算する。
f:id:smijake3:20180720215146p:plain
 {u} {v}から始めて、LCAである頂点 {w}にそれぞれ最後に色 {c_u} {c_v}を通って到達した時の最大コストを計算し、全ての {c_u} {c_v}の組み合わせの内の最大コストを計算すれば解が求まる。

実装

計算量は {O((N+Q) \log N)}

制約が大きいのでPyPyでは厳しかった(TLE)
提出コード (C++14): https://www.codechef.com/viewsolution/19213737

#define N 500005
#define L 21

int n, m, l;

map<int, vector<int> > cc[N];
vector<IP> g[N];
int parent[L][N], depth[N];
// pp[v][k][c1][c2] := vから2^k上の親頂点w間のパスで、端の辺の色がc1, c2の場合の最大コスト
// pll[v][k][i] := vから2^k上の親頂点w間のパスで、頂点vから通る辺の色を持つ
// prr[v][k][i] := vから2^k上の親頂点w間のパスで、頂点wに到達する際に通る辺の色を持つ
int pp[L][N][2][2], pll[L][N][2], prr[L][N][2];
int used[N];

// クエリ内で、最後に通った色と、現在の頂点n0、移動先の親の情報(2^l0個上の親頂点)から
// 次の親頂点での最大コストを計算
void update(vector<P> &e0, int l0, int n0) {
  map<int, int> r;
  repe(ee, e0) {
    int e = ee.first, v0 = ee.second;
    int v2;
    rep(p0, 2) {
      int s0 = pll[l0][n0][p0];
      if(s0 < 0) continue;
      rep(p1, 2) {
        int t0 = prr[l0][n0][p1];
        if(t0 < 0) continue;
        int v1 = pp[l0][n0][p0][p1];
        if(e != -1 && (e != s0 || e == 0 || s0 == 0)) {
          v2 = v0+v1+1;
        } else {
          v2 = v0+v1;
        }
        r[t0] = maxd(r[t0], v2);
      }
    }
  }
  e0.clear();
  repe(ee, r) {
    e0.push_back(P(ee.first, ee.second));
  }
}

// 各クエリに対する計算
// ダブリングでLCAを求めるのに合わせて計算していく
int query(int u, int v) {
  int dd = depth[v] - depth[u];
  if(dd < 0) {
    dd = -dd;
    int tmp = u; u = v; v = tmp;
  }
  // cu (cv) = (c, v): u (v) から始めて、最後にcの色を通って現在の頂点まで移動した時の最大コストv
  vector<P> cu, cv;
  cu.push_back(P(-1, 0)); cv.push_back(P(-1, 0));

  // 両方の現在の頂点の高さ合わせ
  rep(k, l+1) {
    if(dd & 1) {
      update(cv, k, v);
      v = parent[k][v];
    }
    dd >>= 1;
  }
  if(u == v) {
    int r = 0;
    repe(e, cv) {
      r = maxd(r, e.second);
    }
    return r;
  }
 
  // LCA wまで計算
  repr(k, l) {
    if(parent[k][u] != parent[k][v]) {
      update(cu, k, u);
      update(cv, k, v);
      u = parent[k][u];
      v = parent[k][v];
    }
  }
  update(cu, 0, u);
  update(cv, 0, v);

  // LCA wの時点の最大コストを計算
  int r = 0, v2 = 0;
  repe(e0, cu) {
    int eu = e0.first, vu = e0.second;
    repe(e1, cv) {
      int ev = e1.first, vv = e1.second;
      if(eu != ev || eu == 0 || ev == 0) {
        v2 = vu + vv + 1;
      } else {
        v2 = vu + vv;
      }
      r = maxd(r, v2);
    }
  }
  return r;
}
 
int main() {
  scanf("%d %d", &n, &m);
  rep(i, m) {
    int u, v, w; scanf("%d %d %d", &u, &v, &w);
    if(v < u) {
      if(cc[v-1][u-1].size() <= 3) {
        cc[v-1][u-1].push_back(w);
      }
    } else {
      if(cc[u-1][v-1].size() <= 3) {
        cc[u-1][v-1].push_back(w);
      }
    }
  }
  // 色が3つ以上存在する多重辺を1つにする
  // 2色以下はそのまま
  rep(u, n) {
    vector<IP>& g0 = g[u];
    repe(e, cc[u]) {
      int v = e.first;
      vector<int> &s0 = e.second;
      if(s0.size() >= 3) {
        P p0 = P(0, -1);
        g0.push_back(IP(v, p0));
        g[v].push_back(IP(u, p0));
      } else if(s0.size() == 2) {
        int v0 = s0[0], v1 = s0[1];
        P p0 = P(v0, v1);
        g0.push_back(IP(v, p0));
        g[v].push_back(IP(u, p0));
      } else {
        int v0 = s0[0];
        P p0 = P(v0, -1);
        g0.push_back(IP(v, p0));
        g[v].push_back(IP(u, p0));
      }
    }
  }
 
  rep(i, L) {
    rep(j, n) {
      parent[i][j] = -1;
      pp[i][j][0][1] = pp[i][j][1][0] = pp[i][j][1][1] = -INF;
    }
  }
 
  // LCAとdp計算のための前計算
  queue<int> que;
  que.push(0);
  used[0] = 1;
  int dmax = 0;
  while(!que.empty()) {
    int v = que.front(); que.pop();
    rep(i, g[v].size()) {
      IP &e = g[v][i];
      int w = e.first;
      if(used[w]) continue;
      used[w] = 1;
      parent[0][w] = v;
      depth[w] = depth[v] + 1;
      dmax = maxd(dmax, depth[w]);
      P &cc = e.second;
      int e0 = cc.first;

      // 色情報の設定 (1色 or 2色)
      pp[0][w][0][0] = 0;
      pll[0][w][0] = prr[0][w][0] = e0;
      int e1 = cc.second;
      if(e1 != -1) {
        pp[0][w][1][1] = 0;
        pll[0][w][1] = prr[0][w][1] = e1;
      } else {
        pll[0][w][1] = prr[0][w][1] = -1;
      }
      que.push(w);
    }
  }
  l = 1;
  int lev = 1;
  while(lev < dmax) {
    l++;
    lev <<= 1;
  }
 
  // ダブリングによるdpの計算
  rep(k, l-1) {
    int p;
    rep(i, n) {
      pll[k+1][i][0] = pll[k][i][0];
      pll[k+1][i][1] = pll[k][i][1];
      if((p = parent[k][i]) != -1) {
        parent[k+1][i] = parent[k][p];
        prr[k+1][i][0] = prr[k][p][0];
        prr[k+1][i][1] = prr[k][p][1];
        rep(p0, 2) {
          int s0 = pll[k][i][p0];
          if(s0 < 0) continue;
          rep(p1, 2) {
            int t0 = prr[k][i][p1];
            if(t0 < 0) continue;
            int v0 = pp[k][i][p0][p1];
            rep(q0, 2) {
              int s1 = pll[k][p][q0];
              if(s1 < 0) continue;
              rep(q1, 2) {
                int t1 = prr[k][p][q1];
                if(t1 < 0) continue;
                int v1 = pp[k][p][q0][q1];
                if(t0 != s1 || t0 == 0 || s1 == 0) {
                  pp[k+1][i][p0][q1] = maxd(pp[k+1][i][p0][q1], v0 + v1 + 1);
                } else {
                  pp[k+1][i][p0][q1] = maxd(pp[k+1][i][p0][q1], v0 + v1);
                }
              }
            }
          }
        }
      }
    }
  }

  // 各クエリについて計算
  int q; scanf("%d", &q);
  rep(i, q) {
    int u, v;
    scanf("%d %d", &u, &v);
    printf("%d\n", query(u-1, v-1));
  }
  return 0;
}