Submission #6335992

---

Source Code Expand

#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

Judge Result