Official

C - ネットワークの通信コスト / Network Communication Cost Editorial by kyopro_friends


問題文の指示通り、各クエリに対して毎回実際に中継局 A から DFS を行うことで、クエリあたり \(O(N)\) 、全体で \(O(NQ)\) で答えを求めることができます。

python の再帰関数は非常に遅いため、以下の実装例ではキューを用いてDFSを実装しています。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n, q;
  cin >> n >> q;
  vector<int> x(n), y(n);
  for(int i=0; i<n; i++) cin >> x[i] >> y[i];
  vector<vector<int>>g(n);
  for(int i=0; i<n-1; i++){
    int u, v;
    cin >> u >> v;
    u--, v--;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  
  vector<long long> count(n);  // 始点からのパスが含む頂点数
  vector<long long> dist(n);  // 始点からのコスト
  auto dfs = [&](int v){
    // v を始点として DFS する
    count[v] = 1;
    dist[v] = 0;
    queue<array<int,2>>q;
    q.push({v, -1});
    while(q.size() > 0){
      auto[v, pre] = q.front(); q.pop();
      for(auto vv: g[v]){
        if(vv == pre){
          continue;
         }
         count[vv] = count[v] + 1;
         dist[vv] = dist[v] + abs(x[v]-x[vv]) + abs(y[v]-y[vv]);
         q.push({vv, v});
      }
    }
  };
  
  long long s;
  for(int i=0; i<q; i++){
    int t;
    cin >> t;
    if(t == 1){
      int i;
      cin >> i;
      i--;
      x[i] = -x[i];
      y[i] = -y[i];
    }else if(t == 2){
      int w;
      cin >> w;
      s += w;
    }else{
      int a, b;
      cin >> a >> b;
      a--, b--;
      dfs(a);
      cout << dist[b] + s*count[b] << endl;
    }
  }
}

実装例 (Python)

N, Q = map(int, input().split())
XY = [tuple(map(int, input().split())) for _ in range(N)]
G=[[] for _ in range(N)]
for _ in range(N-1):
  U, V = map(int, input().split())
  U -= 1
  V -= 1
  G[U].append(V)
  G[V].append(U)

count = [0]*N  # 始点からのパスが含む頂点数
dist = [0]*N  # 始点からのコスト
def dfs(v):
  # v を始点として DFS する
  count[v] = 1
  dist[v] = 0
  q = [(v, -1)]
  for v, pre in q:
    for vv in G[v]:
      if vv == pre:
        continue
      count[vv] = count[v] + 1
      dist[vv] = dist[v] + abs(XY[v][0]-XY[vv][0]) + abs(XY[v][1]-XY[vv][1])
      q.append((vv, v))

S = 0
for _ in range(Q):
  t, *x = input().split()
  if t == '1':
    i = int(x[0])-1
    XY[i] = (-XY[i][0], -XY[i][1])
  if t == '2':
    S += int(x[0])
  if t == '3':
    A, B = map(int, x)
    A -= 1
    B -= 1
    dfs(A)
    print(dist[B] + S*count[B])

posted:
last update: