So the question is like Given a Tree T, with N nodes. Each node of a tree T has a value associated with it given as A1…AN. You have to perform N-1 queries, and in each query you have to delete the given edge and find the number of unordered pairs (u,v), where |Au−Av | ≤D.
constraints:
|Ai|<=10^6
D<=10
N<=10^5
I am not able to think of a solution that would give AC in time of 1s.
You can do it by adding the edges from last query.Use dsu to add the edges.Removing the (n-1)st edge is same as adding as the tree becomes empty.You can do the quering part by storing the values of the added node and performing binary search .
Can you elaborate the last line more please.
I have no clue how to solve this using binary search, but it’s very simple with policy based data structures.
Do a normal DSU, except maintain an indexed set at each parent, containing all the values of the nodes below the parent. Use normal DSU optimisation techniques to make sure each element gets copied only O(\alpha(n)) times. every time you add a set to another set while merging, check the number of nodes within the good range and add that to our answer. Remember to add the elements only after doing this for all the nodes in the child.
#include <iostream>
#include <bits/stdc++.h>
#include <cmath>
#include <vector>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#define ll long long int
#define mp make_pair
#define pb push_back
#define vi vector<int>
using namespace std;
using namespace __gnu_pbds;
typedef tree<pair<int, int>,null_type, less<pair<int, int>>,rb_tree_tag, tree_order_statistics_node_update> indexed_set;
vector<int> nodeval;
ll ans=0;
int n,d;
struct DSU{
vector<int> parent;
vector<int> sizof;
vector<indexed_set> find;
int n;
DSU(int N){
n=N;
parent.resize(n);
sizof.resize(n);
find.resize(n);
for(int i=0;i<n;i++){
parent[i]=i;
sizof[i]=1;
find[i].insert(mp(nodeval[i], i));
}
}
int get_size(int x){
return sizof[find_set(x)];
}
int find_set(int v){
if (v == parent[v])
return v;
return parent[v] = find_set(parent[v]);
}
void join_sets(int a, int b) {
a = find_set(a);
b = find_set(b);
if (a != b) {
if (sizof[a] < sizof[b]){
swap(a, b);
}
parent[b] = a;
sizof[a] += sizof[b];
for(const auto &val : find[b]){
ans+=find[a].order_of_key(mp(val.first+d, n+5)) - find[a].order_of_key(mp(val.first-d, -1));
}
for(const auto &val : find[b]){
find[a].insert(val);
}
}
}
};
template<typename T>
ostream& operator+(ostream& out, const vector<T> &vec){
for(const auto &x : vec){
out<<x<<" ";
}
out<<"\n";
return out;
}
template<typename T>
ostream& operator*(ostream& out, const vector<T> &vec){
for(const auto &x : vec){
out+x;
}
return out;
}
template<typename T>
istream& operator>>(istream& in, vector<T> &vec){
for(auto &x : vec){
in>>x;
}
return in;
}
void solve(){
cin>>n>>d;
vector<pair<int, int>> edges(n-1);
nodeval.resize(n);
cin>>nodeval;
DSU currtree(n);
for(auto &edge : edges){
cin>>edge.first>>edge.second;
--edge.first;--edge.second;
}
vector<int> order(n-1);
for(int i=n-2;i>=0;--i){
cin>>order[i];
--order[i];
}
vector<ll> fin(n-1);
for(int i=0;i<n-1;i++){
fin[n-2-i]=ans;
currtree.join_sets(edges[order[i]].first, edges[order[i]].second);
}
for(const auto &x : fin){
cout<<x<<'\n';
}
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL);
solve();
}
I used simple vectors for this instead of indexed set, and performed merge operation on them (mergesort) but that approach gets TLE, any idea why?
P.S. I have used union by size.
I’ll need the code:/
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
int n, d;
ll ans = 0;
vector<int> par, p, siz;
vector<pii> edges;
vector<vector<int>> vecs;
ll merge(vector<int>& a, vector<int>& b) {
ll ret = 0;
for (int e : a) {
int id1 = lower_bound(b.begin(), b.end(), e - d) - b.begin();
int id2 = upper_bound(b.begin(), b.end(), e + d) - b.begin();
ret += id2 - id1;
}
vector<int> tmp(a.size() + b.size());
int id = 0, i = 0, j = 0;
while (i < a.size() && j < b.size()) {
if (a[i] <= b[j])
tmp[id++] = a[i++];
else
tmp[id++] = b[j++];
}
while (i < a.size()) tmp[id++] = a[i++];
while (j < b.size()) tmp[id++] = b[j++];
a.resize(tmp.size());
for (i = 0; i < tmp.size(); i++)
a[i] = tmp[i];
b.clear();
return ret;
}
int root(int a) {
if (par[a] == -1) par[a] = a;
int _a = a;
while (a != par[a]) a = par[a];
par[_a] = a;
return a;
}
void add(int a, int b) {
a = root(a);
b = root(b);
if (a == b) return;
if (siz[a] < siz[b]) swap(a, b);
par[b] = a;
siz[a] += siz[b];
ans += merge(vecs[a], vecs[b]);
}
void init() {
edges.resize(n - 1);
p.resize(n - 1);
par.assign(n, -1);
siz.assign(n, 1);
vecs.resize(n, vector<int>(1));
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(NULL);
cout.tie(NULL);
cin >> n >> d;
init();
for (int i = 0; i < n; i++)
cin >> vecs[i][0];
for (int i = 0; i < n - 1; i++) {
int u, v;
cin >> u >> v;
u--, v--;
edges[i] = {u, v};
}
for (int& e : p) cin >> e, e--;
reverse(p.begin(), p.end());
vector<ll> fin(n - 1, 0);
for (int i = 0; i < n - 1; i++) {
pii e = edges[p[i]];
fin[i] = ans;
add(e.first, e.second);
}
reverse(fin.begin(), fin.end());
for (ll& e : fin) cout << e << '\n';
}When you are merging the two vectors, the memory used is the sum of their sizes. You need to do it in O(size of smaller). If i merged one edge to the main tree each time, you would use O(n^2) memory
Although my code is in memory limits, its TLEing. But your point is valid nonetheless.
Actually the same applies for time taken. If b is the smaller and a is the larger, then my code works in O(b\log{a}), whereas merging would take O(a+b)
Thanks!