Submission #1178244

#	Time	Username	Problem	Language	Result	Execution time	Memory
1178244		Andrei_ierdnA	Traffickers (RMI18_traffickers)	C++17	100 / 100	343 ms	16696 KiB

traffickers

#include <iostream>
#include <algorithm>
#include <vector>
#include <random>
#include <chrono>

using namespace std;

using ll = long long;

#define int long long

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MAXN = 30000;
const int MAXK = 50000;
const int MAXQ = 50000;
const int MAXLEN = 20;

int n, k, q;

namespace tree
{
	int root, p[MAXN+2], sz[MAXN+2], depth[MAXN+2];
	int preord[MAXN+2], prepoz[MAXN+2], t;
	int jump[MAXN+2];
	vector<int> nb[MAXN+2];
	
	void dfs(int node)
	{
		preord[++t] = node;
		prepoz[node] = t;
		sz[node] = 1;
		bool bigjump = (depth[node] - depth[jump[node]] == depth[jump[node]] - depth[jump[jump[node]]]);
		for (auto x : nb[node]) {
			if (!p[x]) {
				p[x] = node;
				depth[x] = depth[node] + 1;
				jump[x] = bigjump ? jump[jump[node]] : node;
				dfs(x);
				sz[node] += sz[x];
			}
		}
	}
	void init()
	{
		root = rng() % n + 1;
		// root = 1;
		p[root] = root;
		depth[root] = 1;
		jump[root] = root;
		dfs(root);
	}
	
	int lca(int a, int b)
	{
		if (depth[a] < depth[b]) swap(a, b);
		while (depth[a] > depth[b]) {
			if (depth[jump[a]] >= depth[b]) {
				a = jump[a];
			}
			else {
				a = p[a];
			}
		}
		while (a != b) {
			if (jump[a] != jump[b]) {
				a = jump[a], b = jump[b];
			}
			else {
				a = p[a], b = p[b];
			}
		}
		return a;
	}
}

using namespace tree;

namespace aib
{
	int v[MAXN+2][MAXLEN+1];
	void init()
	{
		for (int i = 1; i <= n; i++) {
			for (int j = 1; j <= MAXLEN; j++) {
				int i2 = i + (i & (-i));
				if (i2 <= n) {
					v[i2][j] += v[i][j];
				}
				int j2 = j + (j & (-j));
				if (j2 <= MAXLEN) {
					v[i][j2] += v[i][j];
				}
			}
		}
	}
	
	void update1(int x, int y, int val)
	{
		while (y <= MAXLEN) {
			v[x][y] += val;
			y += y & (-y);
		}
	}
	void update2(int x, int y, int val)
	{
		while (x <= n) {
			update1(x, y, val);
			x += x & (-x);
		}
	}
	
	int query1(int x, int y)
	{
		int ans = 0;
		while (y) {
			ans += v[x][y];
			y &= y-1;
		}
		return ans;
	}
	int query2(int x, int y)
	{
		int ans = 0;
		while (x) {
			ans += query1(x, y);
			x &= x-1;
		}
		return ans;
	}
}

struct Query
{
	int tip, x, y, lca;
	int t1, t2;
	ll ans;
};
Query qs[MAXK+MAXQ+2];

void solveLen(int len)
{
	for (int i = 1; i <= k+q; i++) {
		if (qs[i].tip <= 2) {
			if (depth[qs[i].x] + depth[qs[i].y] - 2 * depth[qs[i].lca] != len) continue;
			int multi = (qs[i].tip == 1) ? +1 : -1;
			int x = qs[i].x, y = qs[i].y, lca = qs[i].lca;
			int st = 0, dr = len;
			while (x != lca) {
				aib::update2(prepoz[x], st+1, multi);
				aib::update2(prepoz[x]+sz[x], st+1, -multi);
				x = p[x];
				st++;
			}
			aib::update2(prepoz[lca], st+1, multi);
			aib::update2(prepoz[lca]+sz[lca], st+1, -multi);
			while (y != lca) {
				aib::update2(prepoz[y], dr+1, multi);
				aib::update2(prepoz[y]+sz[y], dr+1, -multi);
				y = p[y];
				dr--;
			}
		}
		else {
			int aux = aib::query2(prepoz[qs[i].x], len+1)
			+ aib::query2(prepoz[qs[i].y], len+1)
			- aib::query2(prepoz[qs[i].lca], len+1);
			if (qs[i].lca != tree::root) aux -= aib::query2(prepoz[p[qs[i].lca]], len+1);
			qs[i].ans += 1LL * aux * ((qs[i].t2+1) / (len+1));
			if ((qs[i].t2 + 1) % (len+1)) {
				int amogus = (qs[i].t2 + 1) % (len+1);
				qs[i].ans = qs[i].ans
				+ aib::query2(prepoz[qs[i].x], amogus)
				+ aib::query2(prepoz[qs[i].y], amogus)
				- aib::query2(prepoz[qs[i].lca], amogus);
				if (qs[i].lca != tree::root) qs[i].ans -= aib::query2(prepoz[p[qs[i].lca]], amogus);
			}
			if (qs[i].t1 >= 0) {
				qs[i].ans -= 1LL * aux * ((qs[i].t1+1) / (len+1));
				if ((qs[i].t1 + 1) % (len+1)) {
					int amogus = (qs[i].t1 + 1) % (len+1);
					qs[i].ans = qs[i].ans
					- aib::query2(prepoz[qs[i].x], amogus)
					- aib::query2(prepoz[qs[i].y], amogus)
					+ aib::query2(prepoz[qs[i].lca], amogus);
					if (qs[i].lca != tree::root) qs[i].ans += aib::query2(prepoz[p[qs[i].lca]], amogus);
				}
			}
		}
	}
	for (int i = 1; i <= k+q; i++) {
		if (qs[i].tip <= 2) {
			if (depth[qs[i].x] + depth[qs[i].y] - 2 * depth[qs[i].lca] != len) continue;
			int multi = (qs[i].tip == 1) ? -1 : +1;
			int x = qs[i].x, y = qs[i].y, lca = qs[i].lca;
			int st = 0, dr = len;
			while (x != lca) {
				aib::update2(prepoz[x], st+1, multi);
				aib::update2(prepoz[x]+sz[x], st+1, -multi);
				x = p[x];
				st++;
			}
			aib::update2(prepoz[lca], st+1, multi);
			aib::update2(prepoz[lca]+sz[lca], st+1, -multi);
			while (y != lca) {
				aib::update2(prepoz[y], dr+1, multi);
				aib::update2(prepoz[y]+sz[y], dr+1, -multi);
				y = p[y];
				dr--;
			}
		}
	}
}

signed main()
{
	cin.tie(0)->sync_with_stdio(0);
	cin >> n;
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y;
		tree::nb[x].push_back(y);
		tree::nb[y].push_back(x);
	}
	tree::init();
	cin >> k;
	for (int i = 1; i <= k; i++) {
		qs[i].tip = 1;
		cin >> qs[i].x >> qs[i].y;
		qs[i].lca = tree::lca(qs[i].x, qs[i].y);
	}
	cin >> q;
	for (int i = 1; i <= q; i++) {
		cin >> qs[k+i].tip >> qs[k+i].x >> qs[k+i].y;
		qs[k+i].lca = tree::lca(qs[k+i].x, qs[k+i].y);
		if (qs[k+i].tip == 3) {
			cin >> qs[k+i].t1 >> qs[k+i].t2;
			qs[k+i].t1--;
		}
	}
	for (int len = 0; len < MAXLEN; len++) {
		solveLen(len);
	}
	for (int i = 1; i <= q; i++) {
		if (qs[k+i].tip == 3) {
			cout << qs[k+i].ans << '\n';
		}
	}
	return 0;
}

• Subtask 1: 15 / 15
• Subtask 2: 45 / 45
• Subtask 3: 40 / 40