日々drdrする人のメモ

今日も明日もdrdr

CodeChef July Challenge 2019: Maximum and Minimum

CodeChef July Challenge 2019 の問題: Maximum and Minimum (Code: MXMN)
問題ページ: Contest Page | CodeChef

問題概要

無向グラフ Gに対して関数  f(G, x, y) を定義する。

  •  f(G, x, y) = 頂点 xから頂点 yまでの全てのパスの内での、パスの重みの最小値
    • パスの重みは、パスに含まれる辺の重みの内の最大値と定義

頂点数が N、辺の数が Mの2つの無向グラフ G_1 G_2が与えられる。

以下の値を modulo 998244353 で計算せよ

 \displaystyle S = \sum_{i=1}^{N-1} \sum_{j=i+1}^N f(G_1, i, j) \cdot f(G_2, i, j)

制約

  •  1 \le N \le 10^5
  •  M = 2N
  • 各辺の重み 1 \le w \le 10^8
  •  G_1, G_2は連結グラフ

解法

Heavy-Light Decompositionを用いた上で、DFSを行いながら全頂点ペアの重みを計算する方法で解いた。


まず、各グラフ Gには辺が 2N本存在するが、考えるべき辺は最小全域木上にある (N-1)であることが分かる。

その(大雑把な)理由として、グラフ G最小全域木 Tを考えた時、 Tに含まれない辺 e Tに追加してできる閉路の中で辺 eが最も重みが大きい辺となることから、任意の2点 x, yについて、 Tに含まれない辺を通るパスの重みは Tに含まれる辺のみを通るパス以上の重みになるため、 f(G, x, y) = f(T, x, y)になる、という感じ。

そのため、ここからは2つのグラフ G_1, G_2から最小全域木 T_1, T_2を考える。


各頂点ペア (x, y)ごとに2つの木 T_1, T_2におけるパスの重みを一々計算するのは重くなる。(これを実装すると部分点20がとれる)
そこで、 T_1の各辺 eに対し、 f(T_1, x, y)が辺 eの重みと一致する全てのペア (x, y)に対する T_2内のパスの重みの総和を求めていく方針を考える。


ここで木 T_1, T_2に対し、Kruskal法における辺のマージ順序を元にした二分木を生成する。(生成した木をそれぞれ U_1, U_2とする)
 (N-1)本の辺も頂点とみなし、辺の重みが小さい方から二分木をマージしていき (2N-1)個の頂点からなる二分木を作る。
また、辺が対応する (N-1)個の頂点の重みを、その辺の重みとする。

f:id:smijake3:20190720130848p:plain
生成する二分木のイメージ

この二分木上において、 f(T, x, y)の値は頂点 xに対応する頂点と頂点 yに対応する頂点のLCAの頂点の重みと一致する。(例えば、上の木において f(T, v_2, v_7) v_2の頂点と v_7の頂点のLCAの頂点の重み(=  e_6の重み)と一致する)


解の計算は、木 U_1上でDFSをしながら、木 U_2の上でパスの重みの総和を求めることによって行う。
 U_1上の各頂点では、その左部分木と右部分木に含まれるそれぞれの T_1の頂点集合同士の各ペアについて U_2上でのLCAとなる頂点の重みの総和を求める。そして、その総和に U_1で現在見ている頂点の重み(= 辺 eの重み)を掛けた値を解に足していく。

例えば、以下の図の橙色の頂点では、 U_2上で (v_1, v_3), (v_1, v_4), (v_2, v_3), (v_2, v_4)の各LCAの頂点の重みを求め、その総和を求めた上で辺 eの重みを掛ける。

f:id:smijake3:20190721211309p:plain
DFSの計算のイメージ
橙色の頂点がDFSで現在見ている頂点


問題となるのは、2つの頂点集合 X, Y同士のペアについて U_2上で重みの総和を求める方法であるが、これはHeavy-Light Decompositionを用いて処理する。

 |X| \le |Y| とする。HL分解で U_2の頂点を縮約し、各縮約頂点ごとに、そのheavy-pathに含まれる各頂点について、 Y内の頂点の中でheavy-pathの一番深い頂点とのLCAがその頂点となる回数と、その回数にその頂点の重みを掛けた値を管理するBinary Indexed Treeを持つデータ構造を考える。

f:id:smijake3:20190722212520p:plain
HL分解上で管理するデータのイメージ
緑頂点にかかれている値は管理する重みの総和 (= その頂点の重み × LCAとなる回数)

このデータ構造を用いると、 T_2上の任意の1頂点について、頂点集合 Yに含まれる頂点とのLCAとなる頂点の重みの総和を求められる。(1頂点ごとの計算は O(\log^2 N)で実現できる)
今回は、頂点集合 Xに含まれる各頂点ごとにこの総和を計算していく。

HL分解上で管理する1つのデータ構造を、DFS中の各頂点ごとに一から生成するのは重いため、DFS中は片方の頂点集合をデータ構造に1頂点ずつマージしながらボトムアップに処理していくようにする。(1頂点をデータ構造にマージするのには  O(\log^2 N))


この総和計算とマージ処理は一見すると、全体で O(N^2 \log^2 N)かかるように見えるが、頂点集合が小さい方から頂点集合が大きい方にデータをマージもしくは総和計算するようにすれば、 O(N \log^3 N)で計算することができる。

(参考: "データ構造をマージする一般的なテク" とは? - (iwi) { 反省します - TopCoder部)

また、データ構造1つには空間量  O(N) が必要となる。
単純に考えると、 N個の頂点集合のデータ管理が必要なので空間量  O(N^2)かかりそうに見えるが、 U_1上におけるHL分解を考え、各縮約頂点ごとに一つのデータ構造を用いて管理するようにし、DFSで現在見ている頂点から根頂点までのデータ構造のみを持つ(それ以外は廃棄して再利用する)ようにすることで、空間量  O(N \log N)で計算できるようになる。
(HL分解では、縮約後の木の高さは O(\log N)になることから、必要なデータ構造の数も O(\log N)個になるため)


この解法は全体で、

  • 計算量:  O(N \log^3 N)
  • 空間量:  O(N \log N)

になる。

実装

部分点(20点の方)は木 T_1, T_2上で各点同士の最小重みを計算する。 計算量は  O(N^2)
(部分点) 提出コード: Solution: 25202495 | CodeChef

提出コード(C++14): Solution: 25263215 | CodeChef

#define LV 18
#define N 100003
#define N0 (1 << 17)
#define MOD 998244353

struct Edge {
  int cost;
  int u, v;
  bool operator<(const Edge &other) {
    return cost < other.cost;
  }
};

// union-find data structure
int p[N];
void init(int n) {
  for(int i=0; i<n; ++i) p[i] = i;
}
int root(int x) {
  if(p[x] == x) return x;
  return p[x] = root(p[x]);
}
bool unite(int x, int y) {
  int px = root(x), py = root(y);
  if(px == py) {
    return false;
  }
  if(px < py) {
    p[py] = px;
  } else {
    p[px] = py;
  }
  return true;
}


int n, m;
Edge es[2*N];
// g1: 二分木U1, g2: 二分木U2
vector<int> g1[2*N], g2[2*N];
// costs1[k]: U1の頂点kが持つ重み
// costs2[k]: U2の頂点kが持つ重み
int pc[N], costs1[2*N], costs2[2*N];
int gsz2[2*N];

// heavy-light decomposition
int hs[2*N], prv[2*N], sz1[2*N];
int dfs(int v, int p) {
  int sz = 1, mx = 0;
  int heavy = -1;
  prv[v] = p;
  for(int w : g1[v]) {
    if(w == p) continue;
    int c = dfs(w, v);
    if(mx < c) {
      heavy = w;
      mx = c;
    }
    sz += c;
  }
  hs[v] = heavy;
  return sz1[v] = sz;
}

// dd[k]: k番目の縮約頂点の深さ(今回は未使用)
// sz0[k]: k番目の縮約頂点に含まれる、頂点数
// (lbl[v], pss[v]): 縮約前の頂点において、どの縮約頂点に属し、何番目の頂点になるのかを表す
// ptrs[k]: k番目の縮約頂点について、BITノード配列の何番目から使えばよいかを表す
//   (これにより、BITを1つのデータ構造全体で O(N) で表現する)
// fst[k]: k番目の縮約頂点の、縮約前の頂点らの中の初めの頂点
int dd[N], lbl[2*N], pss[2*N], sz0[N], ptrs[N], gc, fst[N];
// cv[k][i]: k番目の縮約頂点の内、i番目の頂点の(辺の)重み
vector<int> cv[N];
void hld(int s) {
  dfs(s, -1);
  queue<P> que;
  que.push(P(s, 0));
  int k = 0, pcur = 0;
  while(!que.empty()) {
    P &p = que.front();
    int v = p.first, d = p.second;
    que.pop();
    fst[k] = v;
    while(v != -1) {
      pss[v] = cv[k].size();
      cv[k].push_back(costs1[v]);
      lbl[v] = k;
      for(int w : g1[v]) {
        if(hs[v] == w || prv[v] == w) continue;
        que.push(P(w, d+1));
      }
      v = hs[v];
    }
    dd[k] = d;
    sz0[k] = cv[k].size();
    ptrs[k] = pcur;
    pcur += sz0[k];
    ++k;
  }
  gc = k;
}

// U2上のデータを管理するクラス
class Solver {
  // BITの各ノードを管理する構造体
  struct Node {
    // lblはデータの新しさを示すラベル、この値が古いとデータ自体が無効とみなす
    // (これにより、データを一々初期化する必要がなくなる)
    ll sum;
    int lbl, cnt;
  } nds[2*N];
  // 各縮約頂点ごとの情報 (含まれる(T1上の)頂点数、データの新しさを示すラベル)
  int ct1[N], lbs[N];
  // 現在の最新ラベルを表す
  int lb0;

  ll res0; int c0;

public:
  vector<int> vs;

  void init(int n, int gc) {
    for(int i=0; i<gc; ++i) {
      lbs[i] = -1;
    }
    for(int i=0; i<n; ++i) {
      nds[i].lbl = -1;
    }
  }
  void reset(int v) {
    lb0 = v;
    vs.clear();
  }
  // 頂点vと、このデータ構造に含まれる頂点集合との
  // LCAとなる頂点の重みの総和を求める
  ll query(int v) {
    ll res = 0;
    int c0 = 0;
    // HL分解の縮約頂点ごとに処理
    while(v != -1) {
      int l = lbl[v];
      int i = pss[v];

      ll co = cv[l][i];

      // BIT上の取得処理
      Node *b_nds = nds+ptrs[l];
      int k = i;
      int cnt = 0;
      while(k) {
        if(lb0 == b_nds[k].lbl) {
          res += b_nds[k].sum;
          res %= MOD;
          cnt += b_nds[k].cnt;
        }
        k -= k & -k;
      }
      // 位置iよりも下にある全ての頂点の重みは全て cv[l][i] として計算
      ll co1 = (lbs[l] == lb0 ? ct1[l] : 0);
      res += co * (co1 - cnt - c0) % MOD;
      res %= MOD;
      c0 = co1;
      v = prv[fst[l]];
    }
    return res;
  }

  // 1頂点をこのデータ構造にマージする
  void add(int v) {
    vs.push_back(v);
    // HL分解の縮約頂点ごとに処理
    while(v != -1) {
      int l = lbl[v];
      int i = pss[v];

      ll co = cv[l][i];

      // BIT上の加算処理
      Node *b_nds = nds+ptrs[l];
      int n = sz0[l];
      int k = i+1;
      while(k <= n) {
        if(lb0 != b_nds[k].lbl) {
          b_nds[k].lbl = lb0;
          b_nds[k].cnt = 1;
          b_nds[k].sum = co;
        } else {
          b_nds[k].cnt += 1;
          b_nds[k].sum += co;
          b_nds[k].sum %= MOD;
        }
        k += k & -k;
      }
      if(lbs[l] != lb0) {
        lbs[l] = lb0;
        ct1[l] = 1;
      } else {
        ct1[l] += 1;
      }
      v = prv[fst[l]];
    }
  }
} sv[2*LV];

// U1上のDFS処理
ll s_dfs(int v, int lv) {
  if(gsz2[v] == 1) {
    sv[lv].add(v);
    return 0;
  }
  int w0 = g2[v][0];
  int w1 = g2[v][1];
  // 左部分木(頂点集合のサイズが大きい方)について先に計算し、
  // この結果のデータ構造をそのまま使う
  ll r0 = s_dfs(w0, lv);
  // 右部分木(頂点集合のサイズが小さい方)について、
  // 新しいデータ構造で計算していく
  sv[lv+1].reset(w1);
  ll r1 = s_dfs(w1, lv+1);

  ll co = costs2[v];

  // 頂点集合同士の重みの総和を計算
  ll res = (r0 + r1) % MOD;
  ll r2 = 0;
  for(int w : sv[lv+1].vs) {
    r2 += sv[lv].query(w);
    r2 %= MOD;
  }
  res += (r2 * co) % MOD;
  // 左部分木のデータ構造に右部分木に含まれる頂点集合をマージ
  for(int w : sv[lv+1].vs) {
    sv[lv].add(w);
  }
  return res % MOD;
}

int main() {
  cin >> n >> m;
  rep(i, m) {
    int u, v, w;
    cin >> u >> v >> w;
    es[i] = Edge{w, u-1, v-1};
  }
  sort(es, es+m);
  // G1に対するKruskal法を行い、二分木U1をつくる
  init(n);
  rep(i, n) pc[i] = i, sz1[i] = 1;
  int cur = n;
  rep(i, m) {
    Edge &e = es[i];
    int pu = pc[root(e.u)], pv = pc[root(e.v)];
    if(unite(e.u, e.v)) {
      int pp = root(e.u);
      g1[cur].push_back(pu);
      g1[cur].push_back(pv);
      costs1[cur] = e.cost;
      pc[pp] = cur++;
    }
  }
  //assert(cur == 2*n-1);
  hld(2*n-2);

  rep(i, LV) sv[i].init(n, gc);

  rep(i, m) {
    int u, v, w;
    cin >> u >> v >> w;
    es[i] = Edge{w, u-1, v-1};
  }
  sort(es, es+m);
  // G2に対するKruskal法を行い、二分木U2をつくる
  init(n);
  rep(i, n) pc[i] = i, gsz2[i] = 1;
  cur = n;
  rep(i, m) {
    Edge &e = es[i];
    int pu = pc[root(e.u)], pv = pc[root(e.v)];
    if(unite(e.u, e.v)) {
      int pp = root(e.u);
      // 頂点集合が大きい方を左部分木にする
      if(gsz2[pu] < gsz2[pv]) {
        g2[cur].push_back(pv);
        g2[cur].push_back(pu);
      } else {
        g2[cur].push_back(pu);
        g2[cur].push_back(pv);
      }
      costs2[cur] = e.cost;
      gsz2[cur] = gsz2[pu] + gsz2[pv];
      pc[pp] = cur++;
    }
  }
  //assert(cur == 2*n-1);

  // DFSを行い、解を計算
  sv[0].reset(2*n-2);
  ll ans = s_dfs(2*n-2, 0);

  cout << ans << endl;

  return 0;
}