verify/yukicoder_3407.cpp

• View this file on GitHub
• Last update: 2026-05-30 18:08:28+09:00
• Problem: https://yukicoder.me/problems/3407

Depends on

• templates/alias.hpp
• templates/macro.hpp
• AuxiliaryTree (tree/auxiliary_tree.hpp)
• tree/lca.hpp

Code

// competitive-verifier: PROBLEM https://yukicoder.me/problems/3407
#include <bits/stdc++.h>
using namespace std;
#include "templates/alias.hpp"
#include "templates/macro.hpp"

#include "tree/auxiliary_tree.hpp"

int main() {
  ll N;
  in(N);

  vi D(N, 0);

  vvi VG(N);
  vvpii G(N);

  rep(_, N - 1) {
    ll u, v, w;
    in(u, v, w);

    G[u].emplace_back(v, w), G[v].emplace_back(u, w);
    VG[u].emplace_back(v), VG[v].emplace_back(u);
  }

  AuxiliaryTree<ll> AT(VG);

  auto dfs = [&](ll cur, ll par, ll d, const auto &self) -> void {
    D[cur] = d;
    for (auto [nxt, w] : G[cur]) {
      if (nxt == par)
        continue;

      self(nxt, cur, d + w, self);
    }
  };

  dfs(0, -1, 0, dfs);
  ll Q;
  in(Q);

  vvi AG(N);

  while (Q--) {
    ll K;
    in(K);
    vi X(K);
    in(X);

    ll root = AT.build(X, AG), res = 0;

    auto dfs = [&](ll cur, const auto &self) -> void {
      for (auto nxt : AG[cur]) {
        res += D[nxt] - D[cur];

        self(nxt, self);
      }
    };

    dfs(root, dfs);

    out(res);

    for (auto x : X)
      AG[x].clear();
  }
}

#line 1 "verify/yukicoder_3407.cpp"
// competitive-verifier: PROBLEM https://yukicoder.me/problems/3407
#include <bits/stdc++.h>
using namespace std;
#line 7 "templates/alias.hpp"

template <class T> using VC = std::vector<T>;
template <class T>
using rpriority_queue = std::priority_queue<T, VC<T>, std::greater<T>>;
using ll = long long;
using ld = long double;
using pii = std::pair<ll, ll>;
using vi = VC<ll>;
using vvi = VC<vi>;
using vvvi = VC<vvi>;
using vb = VC<bool>;
using vvb = VC<vb>;
using vf = VC<double>;
using vvf = VC<vf>;
using vpii = VC<pii>;
using vvpii = VC<vpii>;
using si = std::set<ll>;
using spii = std::set<pii>;
using mii = std::map<ll, ll>;
const std::string upperlist = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
const std::string lowerlist = "abcdefghijklmnopqrstuvwxyz";
#define mp make_pair
#define dms << " " <<

constexpr int MOD998 = 998244353;
#line 4 "templates/macro.hpp"

// 引数の長さで内容が変わるrep 参考: https://trap.jp/post/1224
#define overload4(a, b, c, d, ...) d
#define _rep(i, n) for (int i = 0; i < (int)(n); i++)
#define REP(i, a, b) for (int i = (int)(a); i < (int)(b); ++i)
#define rep(...) overload4(__VA_ARGS__, REP, _rep)(__VA_ARGS__)
#define _rrep(i, n) for (int i = n - 1; i >= 0; i--)
#define RREP(i, a, b) for (int i = (int)(b - 1); i >= (int)(a); i--)
#define rrep(...) overload4(__VA_ARGS__, RREP, _rrep)(__VA_ARGS__)

#define all(a) (a).begin(), (a).end()

template <typename T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T1, typename T2>
std::istream &operator>>(std::istream &is, std::pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T>
std::istream &operator>>(std::istream &is, std::vector<T> &v) {
  for (T &in : v)
    is >> in;
  return is;
}

template <typename T>
std::ostream &operator<<(std::ostream &os, const std::vector<T> &v) {
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    os << v[i] << (i + 1 == v.size() ? "" : " ");
  }
  return os;
}

// pythonのprintライクな関数 参考:
// https://nyaannyaan.github.io/library/template/inout.hpp
inline void out() { std::cout << std::endl; }
template <typename T, typename... U, char sep = ' '>
void out(const T &t, const U &...u) {
  std::cout << t;

  if (sizeof...(u))
    std::cout << sep;

  out(u...);
}

// cinの短縮関数 参考: https://nyaannyaan.github.io/library/template/inout.hpp
inline void in() {}
template <typename T, class... U> void in(T &t, U &...u) {
  std::cin >> t;
  in(u...);
}

template <typename T> inline T ceil_div(T a, T b) { return (a + b - 1) / b; }
template <typename T> inline T mod_pow(T a, T n, T mod) {
  T res = 1;
  while (n) {
    if (n % 2 != 0) {
      res *= a;
      res %= mod;
    }

    a *= a;
    a %= mod;

    n >>= 1;
  }

  return res;
}
template <typename T> inline T minus_mod(T a, T b) { return ((a % b) + b) % b; }

template <typename T> void apply_vec(std::vector<T> &v, T (*fn)(T)) {
  for (int i = 0; i < v.size(); i++)
    v[i] = fn(v[i]);
}
#line 6 "verify/yukicoder_3407.cpp"

#line 3 "tree/auxiliary_tree.hpp"

#line 2 "tree/lca.hpp"

#line 4 "tree/lca.hpp"

#line 7 "tree/lca.hpp"

template <std::integral T> struct LcaWithDoubling {
  VC<VC<T>> G, P;
  VC<T> D;
  T N, K;

  LcaWithDoubling(VC<VC<T>> G, T st = 0) : G{G}, N{static_cast<int>(G.size())} {
    assert(N > 0);
    assert(st < N);

    K = std::bit_width((unsigned)N);
    P.resize(K, VC<T>(N, -1));
    D.resize(N);

    auto dfs = [&](const auto &self, T cur, T par, T d) -> void {
      P[0][cur] = par;
      D[cur] = d;

      for (auto nxt : G[cur]) {
        if (nxt == par)
          continue;

        self(self, nxt, cur, d + 1);
      }
    };

    dfs(dfs, st, st, 0);

    rep(b, K - 1) rep(i, N) P[b + 1][i] = P[b][P[b][i]];
  }

  T lca(T u, T v) {
    if (D[u] < D[v])
      std::swap(u, v);

    rep(k, K) {
      if ((D[u] - D[v]) >> k & 1) {
        u = P[k][u];
      }
    }

    if (u == v) {
      return u;
    }

    rrep(k, K) {
      if (P[k][u] != P[k][v]) {
        u = P[k][u], v = P[k][v];
      }
    }

    return P[0][u];
  }
};
#line 5 "tree/auxiliary_tree.hpp"

template <std::integral T> struct AuxiliaryTree {
  LcaWithDoubling<T> lca;
  std::vector<int> order;
  int n;

  AuxiliaryTree(std::vector<std::vector<T>> G) : lca{G} {
    n = G.size();
    order.resize(n);

    int cnt = 0;
    auto dfs = [&](int cur, int par, const auto &self) -> void {
      order[cur] = cnt++;
      for (auto nxt : G[cur]) {
        if (nxt == par)
          continue;
        self(nxt, cur, self);
      }
    };

    dfs(0, -1, dfs);
  }

  /**
   * @fn build
   * AuxiliaryTreeを構築します
   * @param V 頂点のリスト 後で補助に使われた頂点も追加される
   */
  T build(std::vector<T> &V, std::vector<std::vector<T>> &G) {
    int m = V.size();

    assert(n == static_cast<int>(G.size()));
    assert(m > 0);

    std::ranges::sort(V, [&](int a, int b) { return order[a] < order[b]; });
    std::vector<T> st = {V[0]};

    for (int i = 0; i < m - 1; i++) {
      T w = lca.lca(V[i], V[i + 1]);

      if (w != V[i]) {
        T last = st.back();
        st.pop_back();
        while (!st.empty() && lca.D[w] < lca.D[st.back()]) {
          int tp = st.back();
          st.pop_back();
          G[tp].emplace_back(last);
          last = tp;
        }
        if (st.empty() || st.back() != w) {
          st.emplace_back(w), V.emplace_back(w);
          G[w] = {last};
        } else {
          G[w].emplace_back(last);
        }
      }

      st.emplace_back(V[i + 1]);
    }

    for (int i = 0; i + 1 < static_cast<int>(st.size()); i++) {
      G[st[i]].emplace_back(st[i + 1]);
    }

    return st[0];
  }
};

/**
 * @file tree/auxiliary_tree.hpp
 * @brief AuxiliaryTree
 * @auther hidehic0
 * @date 2026-05-29
 * verify: https://atcoder.jp/contests/abc340/submissions/76203724
 */
#line 8 "verify/yukicoder_3407.cpp"

int main() {
  ll N;
  in(N);

  vi D(N, 0);

  vvi VG(N);
  vvpii G(N);

  rep(_, N - 1) {
    ll u, v, w;
    in(u, v, w);

    G[u].emplace_back(v, w), G[v].emplace_back(u, w);
    VG[u].emplace_back(v), VG[v].emplace_back(u);
  }

  AuxiliaryTree<ll> AT(VG);

  auto dfs = [&](ll cur, ll par, ll d, const auto &self) -> void {
    D[cur] = d;
    for (auto [nxt, w] : G[cur]) {
      if (nxt == par)
        continue;

      self(nxt, cur, d + w, self);
    }
  };

  dfs(0, -1, 0, dfs);
  ll Q;
  in(Q);

  vvi AG(N);

  while (Q--) {
    ll K;
    in(K);
    vi X(K);
    in(X);

    ll root = AT.build(X, AG), res = 0;

    auto dfs = [&](ll cur, const auto &self) -> void {
      for (auto nxt : AG[cur]) {
        res += D[nxt] - D[cur];

        self(nxt, self);
      }
    };

    dfs(root, dfs);

    out(res);

    for (auto x : X)
      AG[x].clear();
  }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	challenge01	AC	136 ms	50 MB
g++	sample01	AC	2 ms	3 MB
g++	test_1_01	AC	3 ms	4 MB
g++	test_1_02	AC	3 ms	3 MB
g++	test_1_03	AC	3 ms	3 MB
g++	test_1_04	AC	4 ms	4 MB
g++	test_1_05	AC	3 ms	3 MB
g++	test_2_01	AC	212 ms	49 MB
g++	test_2_02	AC	284 ms	49 MB
g++	test_2_03	AC	194 ms	49 MB
g++	test_2_04	AC	228 ms	49 MB
g++	test_2_05	AC	231 ms	49 MB
g++	test_3_01	AC	194 ms	49 MB
g++	test_3_02	AC	201 ms	49 MB
g++	test_3_03	AC	208 ms	49 MB
g++	test_3_04	AC	208 ms	49 MB
g++	test_3_05	AC	203 ms	49 MB
g++	test_4_01	AC	256 ms	49 MB
g++	test_4_02	AC	237 ms	49 MB
g++	test_4_03	AC	227 ms	49 MB
g++	test_4_04	AC	230 ms	49 MB
g++	test_4_05	AC	306 ms	49 MB
g++	test_5_01	AC	201 ms	49 MB
g++	test_5_02	AC	186 ms	49 MB
g++	test_5_03	AC	190 ms	49 MB
g++	test_5_04	AC	198 ms	49 MB
g++	test_5_05	AC	200 ms	49 MB
g++	test_6_01	AC	258 ms	49 MB
g++	test_6_02	AC	256 ms	49 MB
g++	test_6_03	AC	270 ms	49 MB

Back to top page