Submission #674053

#	Time	Username	Problem	Language	Result	Execution time	Memory
674053		ghostwriter	Construction of Highway (JOI18_construction)	C++11	100 / 100	279 ms	21184 KiB

construction

// #pragma GCC optimize ("Ofast")
// #pragma GCC target ("avx2")
#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL
    #include <debug.h>
#else
    #define debug(...)
#endif
#define ft front
#define bk back
#define st first
#define nd second
#define ins insert
#define ers erase
#define pb push_back
#define pf push_front
#define _pb pop_back
#define _pf pop_front
#define lb lower_bound
#define ub upper_bound
#define mtp make_tuple
#define bg begin
#define ed end
#define all(x) (x).bg(), (x).ed()
#define sz(x) (int)(x).size()
// #define int long long
typedef long long ll; typedef unsigned long long ull;
typedef double db; typedef long double ldb;
typedef pair<int, int> pi; typedef pair<ll, ll> pll;
typedef vector<int> vi; typedef vector<ll> vll; typedef vector<pi> vpi; typedef vector<pll> vpll;
typedef string str;
#define FOR(i, l, r) for (int i = (l); i <= (r); ++i)
#define FOS(i, r, l) for (int i = (r); i >= (l); --i)
#define FRN(i, n) for (int i = 0; i < (n); ++i)
#define FSN(i, n) for (int i = (n) - 1; i >= 0; --i)
#define EACH(i, x) for (auto &i : (x))
#define WHILE while
template<typename T> T gcd(T a, T b) { T d2 = (a | b) & -(a | b); a /= d2; b /= d2; WHILE(b) { a = a % b; swap(a, b); } return a * d2; }
template<typename T> T lcm(T a, T b) { return a / gcd(a, b) * b; }
void _assert(bool statement) { if (statement) return; cerr << "\n>> Assertion failed!\n"; exit(0); }
void _assert(bool statement, const str &message) { if (statement) return; cerr << "\n>> Assertion failed: " << message << '\n'; exit(0); }
void _error(const str &message) { cerr << "\n>> Error: " << message << '\n'; exit(0); }
#define file "TEST"
mt19937 rd(chrono::steady_clock::now().time_since_epoch().count());
ll rand(ll l, ll r) { return uniform_int_distribution<ll>(l, r)(rd); }
/*
----------------------------------------------------------------
    END OF TEMPLATE
----------------------------------------------------------------
    Tran The Bao - ghostwriter
    Training for VOI23 gold medal
----------------------------------------------------------------
    GOAT
----------------------------------------------------------------
*/
const int N = 1e5 + 5;
int n, f[N], c[N], s[N], h[N], p[N], heavy[N], pos[N], head[N], num[N], timer = 1;
vi adj[N];
vpi a[N], hist;
pi e[N];
void input(int test_id) {
	cin >> n;
	FOR(i, 1, n) cin >> c[i];
	FOR(i, 1, n - 1) {
		cin >> e[i].st >> e[i].nd;
		adj[e[i].st].pb(e[i].nd);
	}
}
void compress(int c[]) {
	vi v;
	FOR(i, 1, n) v.pb(c[i]);
	sort(all(v));
	FOR(i, 1, n) c[i] = lb(all(v), c[i]) - v.bg() + 1;
}
void upd(int pos, int v) { hist.pb({pos, v}); for (int i = pos; i <= n; i += (i & -i)) f[i] += v; }
int get(int pos) { int rs = 0; for (int i = pos; i > 0; i -= (i & -i)) rs += f[i]; return rs; }
void clear() { EACH(j, hist) { int pos = j.st, v = j.nd; for (int i = pos; i <= n; i += (i & -i)) f[i] -= v; } hist.clear(); }
void dfs(int u) {
	EACH(v, adj[u]) {
		h[v] = h[u] + 1;
		p[v] = u;
		dfs(v);
		s[u] += s[v];
		if (!heavy[u]) heavy[u] = v;
		else if (s[heavy[u]] < s[v]) heavy[v] = u;
	}
}
void hld(int u) {
	pos[u] = timer;
	if (a[pos[u]].empty()) head[pos[u]] = u;
	a[pos[u]].pb({c[u], 1});
	num[u] = sz(a[pos[u]]);
	if (heavy[u]) hld(heavy[u]);
	EACH(v, adj[u]) {
		if (v == heavy[u]) continue;
		++timer;
		hld(v);
	}
}
ll get(int u, int v) {
	ll rs = 0;
	WHILE(u) {
		vpi a1;
		int cpos = pos[u], cnum = num[u];
		FSN(j, sz(a[cpos])) {
			if (!cnum) break;
			int tk = min(cnum, a[cpos][j].nd);
			bool ok = 1;
			if (tk == a[cpos][j].nd) a[cpos]._pb();
			else {
				a[cpos][j].nd -= tk;
				ok = 0;
			}
			cnum -= tk;
			a1.pb({a[cpos][j].st, tk});
			if (!ok) break;
		}
		reverse(all(a1));
		EACH(i, a1) {
			rs += 1LL * i.nd * get(i.st - 1);
			upd(i.st, i.nd);
		}
		a[cpos].pb({v, num[u]});
		u = p[head[cpos]];
	}
	clear();
	return rs;
}
void solve(int test_id) {
	compress(c);
	dfs(1);
	hld(1);
	FOR(i, 1, timer) reverse(all(a[i]));
	FOR(i, 1, n - 1) {
		int u = e[i].st, v = e[i].nd;
		cout << get(u, c[v]) << '\n';
	}
}
void reinit(int test_id) {

}
signed main() {
    ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    // freopen(file".inp", "r", stdin);
    // freopen(file".out", "w", stdout);
    int test_num = 1;
    // cin >> test_num; // comment if the problem does not requires multitest
    FOR(i, 1, test_num) {
        input(i); // input in noninteractive problems for case #i
        solve(i); // main function to solve case #i
        reinit(i); // reinit global data to default used in case #i
    }
    #ifdef LOCAL
        cerr << "\nTime: " << setprecision(5) << fixed << (ldb)clock() / CLOCKS_PER_SEC << "ms.\n";
    #endif
    return 0;
}
/*
----------------------------------------------------------------
From Benq:
    stuff you should look for
        * int overflow, array bounds
        * special cases (n=1?)
        * do smth instead of nothing and stay organized
        * WRITE STUFF DOWN
        * DON'T GET STUCK ON ONE APPROACH
----------------------------------------------------------------
*/

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