Submission #1001183

#	Time	Username	Problem	Language	Result	Execution time	Memory
1001183		aykhn	Construction of Highway (JOI18_construction)	C++17	100 / 100	639 ms	47564 KiB

construction

#include <bits/stdc++.h>

using namespace std;

#define int long long
#define inf 0x3F3F3F3F3F3F3F3F

struct BIT
{
	int n;
	vector<int> ft;
	void init(int N)
	{
		n = N + 10;
		ft.assign(n + 10, 0);
	}
	void add(int pos, int val)
	{
		for (pos = pos + 5; pos <= n; pos += pos & -pos) ft[pos] += val;
	}
	int get(int pos, int res = 0)
	{
		for (pos = pos + 5; pos > 0; pos -= pos & -pos) res += ft[pos];
		return res;
	}
};

struct SegTree
{
	int sz;
	vector<array<int, 2>> st;
	vector<int> lz;
	void init(int n)
	{
		st.assign((n + 1) << 2, {-inf, inf});
		lz.assign((n + 1) << 2, -1);
	}
	void relax(int l, int r, int x)
	{
		if (lz[x] == -1) return;
		st[x][0] = st[x][1] = lz[x];
		if (l == r)
		{
			lz[x] = -1;
			return;
		}
		lz[2*x] = lz[2*x + 1] = lz[x];
		lz[x] = -1;
	}
	array<int, 2> combine(array<int, 2> a, array<int, 2> b)
	{
		return {max(a[0], b[0]), min(a[1], b[1])};
	}
	void make(int l, int r, int x, int lx, int rx, int val)
	{
		if (lx > rx) return;
		relax(l, r, x);
		if (l > rx || r < lx) return;
		if (l >= lx && r <= rx)
		{
			lz[x] = val;
			relax(l, r, x);
			return;
		}
		int mid = (l + r) >> 1;
		make(l, mid, 2*x, lx, rx, val);
		make(mid + 1, r, 2*x + 1, lx, rx, val);
		st[x] = combine(st[2*x], st[2*x + 1]);
	}
	array<int, 2> get(int l, int r, int x, int lx, int rx)
	{
		if (lx > rx) return {-inf, inf};
		if (l > rx || r < lx) return {-inf, inf};
		relax(l, r, x);
		if (l >= lx && r <= rx) return st[x];
		int mid = (l + r) >> 1;
		return combine(get(l, mid, 2*x, lx, rx), get(mid + 1, r, 2*x + 1, lx, rx));
	}
};

const int MXN = 1e5 + 5;
const int LOG = 18;

int n;
vector<int> adj[MXN];
int sz[MXN], p[LOG][MXN], dep[MXN], head[MXN], col[MXN];
int in[MXN], out[MXN], tim;

void _init(int a)
{
	sz[a] = 1;
	int mx = 0;
	for (int i = 0; i < adj[a].size(); i++)
	{
		dep[adj[a][i]] = dep[a] + 1;
		_init(adj[a][i]);
		sz[a] += sz[adj[a][i]];
		if (sz[adj[a][i]] > sz[adj[a][mx]]) mx = i;
	}
	if (!adj[a].empty()) swap(adj[a][mx], adj[a][0]);
}

void dfs(int a)
{
	in[a] = ++tim;
	for (int i = 0; i < adj[a].size(); i++)
	{
		if (!i) head[adj[a][i]] = head[a];
		else head[adj[a][i]] = adj[a][i];
		dfs(adj[a][i]);
	}
	out[a] = tim;
}

BIT ft;
SegTree st;

void upd(int a, int val)
{
	while (a != 1)
	{
		if (head[a] == a)
		{
			col[a] = val;
			a = p[0][a];
		}
		else
		{
			st.make(1, n, 1, in[head[a]] + 1, in[a], val);
			a = head[a];
		}
	}
	col[1] = val;
}
vector<array<int, 2>> tmp;
int handle(int l, int r)
{
	int ans = 0;
	while (r >= l)
	{
		int lx = l, rx = r;
		int val = st.get(1, n, 1, r, r)[0];
		while (lx < rx)
		{
			int mid = (lx + rx) >> 1;
			array<int, 2> x = st.get(1, n, 1, mid, r);
			if (x[0] == x[1]) rx = mid;
			else lx = mid + 1;
		}
		ans += ft.get(val - 1) * (r - lx + 1);
		tmp.push_back({val, r - lx + 1});
		ft.add(val, r - lx + 1);
		r = lx - 1;
	}
	return ans;
}
int getans(int a)
{
	int res = 0;
	while (a != 1)
	{
		if (head[a] == a)
		{
			res += ft.get(col[a] - 1);
			ft.add(col[a], 1);
			tmp.push_back({col[a], 1});
			a = p[0][a];
		}
		else
		{
			res += handle(in[head[a]] + 1, in[a]);
			a = head[a];
		}
	}
	res += ft.get(col[1] - 1);
	while (!tmp.empty())
	{
		ft.add(tmp.back()[0], -tmp.back()[1]);
		tmp.pop_back();
	}
	return res;
}

vector<int> mp;

signed main()
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> n;
	st.init(n);
	for (int i = 1; i <= n; i++) cin >> col[i], mp.push_back(col[i]);
	sort(mp.begin(), mp.end());
	mp.resize(unique(mp.begin(), mp.end()) - mp.begin());
	for (int i = 1; i <= n; i++) col[i] = lower_bound(mp.begin(), mp.end(), col[i]) - mp.begin() + 1;
	ft.init(mp.size());
	vector<array<int, 2>> e;
	head[1] = 1;
	for (int i = 1; i < n; i++)
	{
		int u, v;
		cin >> u >> v;
		e.push_back({u, v});
		p[0][v] = u;
		adj[u].push_back(v);
	}
	p[0][1] = 1;
	for (int i = 1; i < LOG; i++) for (int j = 1; j <= n; j++) p[i][j] = p[i - 1][p[i - 1][j]];
	_init(1);
	dfs(1);
	st.make(1, n, 1, 1, 1, col[1]);
	for (int i = 0; i < e.size(); i++)
	{
		cout << getans(e[i][0]) << '\n';
		upd(e[i][1], col[e[i][1]]);
	}
}

Compilation message (stderr)

construction.cpp: In function 'void _init(long long int)':
construction.cpp:93:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   93 |  for (int i = 0; i < adj[a].size(); i++)
      |                  ~~^~~~~~~~~~~~~~~
construction.cpp: In function 'void dfs(long long int)':
construction.cpp:106:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  106 |  for (int i = 0; i < adj[a].size(); i++)
      |                  ~~^~~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:212:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 2> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  212 |  for (int i = 0; i < e.size(); i++)
      |                  ~~^~~~~~~~~~

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84