Submission #6335992


Source Code Expand

Copy
#include <bits/stdc++.h>

#define rep(i, a, b) for (ll i = (a); i < (b); i++)
typedef uint64_t ull;
typedef int64_t ll;
typedef std::pair<ll, ll> PLL;
using namespace std;

// 各ノードは、0-indexed
class LCA {
  private:
    int N; // ノード数
    const int NONE = -1;
    std::vector<std::vector<int>> G;
    int root; // 根ノードの番号
    vector<vector<int>> parent; // parent[k][n]: ノードnから2^k回たどって到達する頂点
    vector<int> depth;          // 根からの深さ
    void dfs(int v, int p, int d) {
      parent[0][v] = p;
      depth[v] = d;
      for (auto c:  G[v]) {
        if (c == p)
          continue;
        dfs(c, v, d+1);
      }
    }
  public:
    explicit LCA(int N, int root) : N(N), G(N, vector<int>(0)), parent(63, vector<int>(N)), depth(N), root(root) {}
    void addEdge(int node1, int node2) {
      G[node1].push_back(node2);
      G[node2].push_back(node1);
    }

    void build(){
      // parent[0]とdepthを初期化
      dfs(root, NONE, 0);
      // parentを初期化
      for (int k = 0; k + 1 < 63; k++) {
        for (int v = 0; v < N; v++) {
          if (parent[k][v] == NONE)
            parent[k + 1][v] = NONE;
          else
            parent[k + 1][v] = parent[k][parent[k][v]];
        }
      }
    }

    int query(int node1, int node2) {
      if (depth[node1] > depth[node2])
        swap(node1, node2);
      for (int k = 0; k < 63; k++) {
        if ((depth[node2] - depth[node1]) >> k & 1) {
          node2 = parent[k][node2];
        }
      }
      if (node1 == node2)
        return node1;
      for (int k = 62; k >= 0; k--) {
        if (parent[k][node1] != parent[k][node2]) {
          node1 = parent[k][node1];
          node2 = parent[k][node2];
        }
      }
      return parent[0][node1];
    }

    int numEdge(int node1, int node2) {
      int lca = query(node1, node2);
      return depth[node1] + depth[node2] - 2 * depth[lca];
    }
};

struct Query {
  ll x, y, color, length;
};

struct Edge {
  ll nxt, color, length;
};

ll N, Q;
vector<vector<ll>> whatYouNeed; // whatYouNeed[node]: nodeが事前に知っておくべき色の一覧
vector<ll> sumAllDistance; // sumAllDistance[node]: rootからnodeまでの辺の長さの合計
map<PLL,ll> sumDistance; // sumDistance[node, c]: rootからnodeまでの色cの辺の長さの合計
map<PLL,ll> numEdge; // numEdge[node, c]: rootからnodeまでの色cの辺の数
vector<vector<Edge>> edges;

ll sad; // sad[c]: 現在のnodeにおける辺の長さの合計
vector<ll> sd; // sd[c]: 現在のnodeにおける色cの辺の長さの合計
vector<ll> nd; // nd[c]: 現在のnodeにおける色cの辺の数
void dfs(int node, int par) {
  // メモしておく
  sumAllDistance[node] = sad;
  for (auto col : whatYouNeed[node]) {
    sumDistance[PLL(node, col)] = sd[col];
    numEdge[PLL(node, col)] = nd[col];
  } 

  // 子どもを展開
  for (auto e : edges[node]) {
    if (e.nxt == par)
      continue;
    sad += e.length;
    sd[e.color]+=e.length;
    nd[e.color]++;
    dfs(e.nxt, node);
    nd[e.color]--;
    sd[e.color]-=e.length;
    sad -= e.length;
  }
}

signed main() {
  cin>>N>>Q;
  LCA lca(N, 0);
  vector<Query> queries;
  whatYouNeed.resize(N);
  sumAllDistance.resize(N);
  //sumDistance.resize(N);
  //numEdge.resize(N);
  edges.resize(N);
  sd.resize(N,0);
  nd.resize(N,0);

  rep(i,0,N-1){
    ll a,b,c,d;
    cin>>a>>b>>c>>d;
    a--;
    b--;
    lca.addEdge(a,b);
    edges[a].push_back(Edge{b,c,d});
    edges[b].push_back(Edge{a,c,d});
  }
  lca.build();
  rep(i,0,Q){
    ll a,b,c,d;
    cin>>c>>d>>a>>b;
    a--;
    b--;
    queries.push_back(Query{a,b,c,d});
    whatYouNeed[a].push_back(c);
    whatYouNeed[b].push_back(c);
    whatYouNeed[lca.query(a,b)].push_back(c);
  }
  dfs(0, -1);
  for (auto q: queries) {
    ll a = q.x, b = q.y, c = q.color, d = q.length;

    // 変更前の距離
    ll ans = sumAllDistance[a] + sumAllDistance[b] - 2 * sumAllDistance[lca.query(a,b)];

    // 色cの付いた辺を切り離す
    ans -= sumDistance[PLL(a, c)] + sumDistance[PLL(b, c)] - 2 * sumDistance[PLL(lca.query(a,b),c)];

    // 色cの付いた新たな辺を加える
    ans += d * (numEdge[PLL(a, c)] + numEdge[PLL(b, c)] - 2 * numEdge[PLL(lca.query(a,b),c)]);

    cout<<ans<<endl;
  }
  return 0;
}

Submission Info

Submission Time
Task F - Colorful Tree
User bobuhiro11
Language C++14 (GCC 5.4.1)
Score 600
Code Size 4487 Byte
Status
Exec Time 1529 ms
Memory 86124 KB

Test Cases

Set Name Score / Max Score Test Cases
Sample 0 / 0 a01
All 600 / 600 a01, b02, b03, b04, b05, b06, b07, b08, b09, b10, b11, b12
Case Name Status Exec Time Memory
a01 1 ms 256 KB
b02 1 ms 256 KB
b03 263 ms 7408 KB
b04 1170 ms 72428 KB
b05 1334 ms 84076 KB
b06 1210 ms 62956 KB
b07 1265 ms 66412 KB
b08 1353 ms 86124 KB
b09 1305 ms 83304 KB
b10 1529 ms 79852 KB
b11 1509 ms 77292 KB
b12 1487 ms 79596 KB