Codechef July Long Challenge 2021 Solution | Chef and Pairs solution in C++ | AskTheCode
Problem:
You are given a tree (connected, undirected, acyclic graph) consisting of N nodes. Based on this tree, you have to answer Q queries.
Each query is of the form:
K D V1 V2 ⋯ VK - output the number of pairs (i, j),1≤i<j≤K, such that the shortest path between nodes Vi and Vj in the tree has D edges.
Input:
The first line contains an integer T, the number of test cases. Then the test cases follow.
The first line of each test case contains two integers, N and Q.
N−1 lines follow. Each line consists of two integers u and v, indicating that there is an edge between nodes u and v in the tree.
Q lines follow. Each line describes a query in the format given above.
Output:
For each query, output the answer on a new line.
Sample Input:
1
5 2
1 2
1 3
2 4
4 5
3 2 2 3 5
2 4 1 3Sample Output:
2
0EXPLANATION:
In the first query, the pairs of nodes (2,3) and (2,5) have distance 2.
In the second query, there is no pair with distance 4.
Code:
#include <bits/stdc++.h>
using namespace std;
#define _USE_MATH_DEFINES
struct rooted_tree {
vector<vector<int>> adj;
vector<int> parent, depth, subtree, start, preorder, in, out;
const vector<int>& operator [] (int i) const { return adj[i]; }
rooted_tree(const vector<vector<int>>& adj_list, int root): adj(adj_list),
parent(adj.size(), -1), depth(adj.size()), subtree(adj.size()), start(adj.size()),
in(adj.size()), out(adj.size()) {
preorder.reserve(adj.size());
build(root, -1, 0);
}
rooted_tree(vector<vector<int>>&& adj_list, int root): adj(move(adj_list)),
parent(adj.size(), -1), depth(adj.size()), subtree(adj.size()), start(adj.size()),
in(adj.size()), out(adj.size()) {
preorder.reserve(adj.size());
build(root, -1, 0);
}
bool is_ancestor_of(int anc, int v, bool strict = true) const {
if (strict) return in[anc] < in[v] && out[anc] > out[v];
else return in[anc] <= in[v] && out[anc] >= out[v];
}
private:
int build(int u, int par, int idx) {
in[u] = idx++;
start[u] = (int)preorder.size();
preorder.push_back(u);
parent[u] = par;
subtree[u] = 1;
if (par != -1) {
depth[u] = depth[par] + 1;
vector<int>::iterator it = find(adj[u].begin(), adj[u].end(), par);
if (it != adj[u].end()) {
adj[u].erase(it);
}
}
for (int v : adj[u]) {
idx = build(v, u, idx);
subtree[u] += subtree[v];
}
out[u] = idx++;
return idx;
}
};
struct lca_binary_jumping : rooted_tree {
int L;
vector<vector<int>> jump;
lca_binary_jumping(const vector<vector<int>>& adj_list, int root):
rooted_tree(adj_list, root) {
build();
}
lca_binary_jumping(vector<vector<int>>&& adj_list, int root):
rooted_tree(move(adj_list), root) {
build();
}
int lca(int a, int b) const {
if (depth[a] < depth[b]) std::swap(a, b);
for (int j = L - 1; j >= 0; j--) {
if (depth[a] - (1 << j) >= depth[b]) {
a = jump[j][a];
}
}
if (a == b) return a;
for (int j = L - 1; j >= 0; j--) {
if (jump[j][a] != jump[j][b]) {
a = jump[j][a];
b = jump[j][b];
}
}
return parent[a];
}
int distance(int a, int b) const {
return depth[a] + depth[b] - 2 * depth[lca(a, b)];
}
private:
void build() {
L = 32 - __builtin_clz((int)adj.size());
jump.resize(L, vector<int>(adj.size(), -1));
for (int u : preorder) {
jump[0][u] = parent[u];
for (int j = 1; j < L && jump[j-1][u] != -1; j++) {
jump[j][u] = jump[j-1][jump[j-1][u]];
}
}
}
};
struct compressed_tree {
vector<int> vertices;
map<int, int> remap;
vector<vector<pair<int, int>>> adj;
vector<pair<int, int>> parent;
compressed_tree(const lca_binary_jumping& tree, const vector<int>& verts) {
if (verts.empty()) throw invalid_argument("vertices of compressed tree must not be empty");
vector<pair<int, int>> order;
for (int v : verts) {
order.emplace_back(tree.in[v], v);
}
sort(begin(order), end(order));
for (size_t i = 1; i < verts.size(); i++) {
int lca = tree.lca(order[i].second, order[i - 1].second);
order.emplace_back(tree.in[lca], lca);
}
sort(begin(order), end(order));
int n = (int)(unique(begin(order), end(order)) - begin(order));
vertices.resize(n);
adj.resize(n);
parent.resize(n);
parent[0] = pair(-1, 0);
for (int i = 0; i < n; i++) {
vertices[i] = order[i].second;
remap[vertices[i]] = i;
if (i > 0) {
int lca = tree.lca(vertices[i], vertices[i - 1]);
int dist = tree.distance(vertices[i], lca);
parent[i] = pair(remap[lca], dist);
adj[parent[i].first].emplace_back(i, dist);
}
}
}
const vector<pair<int, int>>& operator [] (int i) const { return adj[i]; }
};
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
template <typename T>
using ordered_set = __gnu_pbds::tree<T,
__gnu_pbds::null_type,
less<T>,
__gnu_pbds::rb_tree_tag,
__gnu_pbds::tree_order_statistics_node_update>;
using ll = long long;
using ld = long double;
using pt = complex<ld>;
constexpr char nl = '\n';
constexpr int INF = 0x3f3f3f3f;
constexpr ll INFLL = 0x3f3f3f3f3f3f3f3f;
constexpr int MOD = 998244353;
constexpr ld EPS = 1e-9L;
random_device _rd; mt19937 rng(_rd());
tuple<ll, map<int, int>, int> solve(const compressed_tree& tree, const vector<bool>& mark, int u, int want) {
ll res = 0;
map<int, int> dists;
if (mark[u]) dists[0] = 1;
int offset = 0;
for (auto [v, c] : tree[u]) {
auto [add, subdists, suboff] = solve(tree, mark, v, want);
res += add;
suboff += c;
if (size(dists) < size(subdists)) {
swap(dists, subdists);
swap(offset, suboff);
}
for (auto [d, cnt] : subdists) {
int actual = d + suboff; // actual distance
if (dists.count(want - actual - offset)) {
res += (ll)dists[want - actual - offset] * cnt;
}
}
for (auto [d, cnt] : subdists) {
dists[d + suboff - offset] += cnt;
}
}
return tuple(res, move(dists), offset);
}
int main() {
cin.tie(0)->sync_with_stdio(0);
cout << fixed << setprecision(10);
#if defined(ONLINE_JUDGE) && defined(FILENAME)
freopen(FILENAME ".in", "r", stdin);
freopen(FILENAME ".out", "w", stdout);
#endif
int T;
cin >> T;
while(T--) {
int n, m;
cin >> n >> m;
vector<vector<int>> adj(n+1);
for(int i=1; i<n; i++) {
int a, b;
cin >> a >> b;
adj[a].push_back(b);
adj[b].push_back(a);
}
lca_binary_jumping tree(move(adj), 1);
while(m--) {
int k, d;
cin >> k >> d;
vector<int> verts(k);
for(int i=0; i<k; i++) {
cin >> verts[i];
}
compressed_tree graph(tree, verts);
vector<bool> mark(graph.adj.size());
for(int v : verts) {
mark[graph.remap[v]] = true;
}
cout << get<0>(solve(graph, mark, 0, d)) << nl;
}
}
return 0;
}
Commentaires