Submission #943304

#	Time	Username	Problem	Language	Result	Execution time	Memory
943304		lovrot	Construction of Highway (JOI18_construction)	C++17	100 / 100	251 ms	27688 KiB

construction

#include <cstdio> 
#include <set>
#include <vector> 
#include <algorithm> 
#include <stack> 

#define X first
#define Y second
#define PB push_back
#define MP make_pair

using namespace std; 

typedef long long ll; 
typedef pair<int, int> pii;

const int LOG = 17;
const int N = 1 << LOG; 

int A[N], B[N], C[N];

int DEP[N], SIZ[N], CH[N], chi, H[N], P[N];
vector<int> G[N], CMP;

int dfs(int u) {
	++SIZ[u];

	int mx = 0, mxi = 0; 
	for(int v : G[u]) {
		DEP[v] = DEP[u] + 1;
		int res = dfs(v);

		SIZ[u] += res;
		if(res > mx) {
			mx = res;
			mxi = v;
		}
	}

	CH[u] = !mxi ? ++chi : CH[mxi]; 
	H[CH[u]] = u;
	return SIZ[u];
}

int F[N];
stack<pii> U;

void add(int x, int val, bool f = 1) {
	if(f) U.push({x, val});
	for(; x < N; x += x & -x) F[x] += val;
}

int sum(int x) {
	int ret = 0;
	for(; x; x -= x & -x) ret += F[x];
	return ret;
}

void undo() {
	for(; !U.empty(); U.pop()) add(U.top().X, -U.top().Y, 0); 
}

set<pii> S[N]; // {DEP, VAL}

ll solve(int u, int x) { // NE ZABORAVI parent mora bit dodan u set od lanca (1 !!!)
	if(!u) {
		undo();
		return 0;
	}
	
//	printf("entry %d\n", u);

	int ch = CH[u];
	auto it = S[ch].lower_bound(MP(DEP[u], -1));

	ll ret = 0;
	for(bool f = (*it).X == DEP[u]; 1; f = 1) {
		int dep = (*it).X, val = (*it).Y;
		int siz = (f ? dep : DEP[u]) - (it == S[ch].begin() ? DEP[H[ch]] - 1 : (*prev(it)).X);

//		printf("%d| %d %d\n", it == S[ch].begin(), siz, sum(val - 1));
		ret += (ll) siz * sum(val - 1);
		add(val, siz);
		
		auto nxt = it == S[ch].begin() ? S[ch].end() : prev(it);

		if(f) S[ch].erase(it);
	
		if(nxt == S[ch].end()) break;
		it = nxt;
	}
	
	S[ch].insert({DEP[u], x});
	return ret + solve(P[H[ch]], x);
}

int main() {
	int n;
	scanf("%d", &n);
	for(int i = 1; i <= n; ++i) {
		scanf("%d", C + i); 
		CMP.PB(C[i]);
	}

	sort(CMP.begin(), CMP.end());
	CMP.erase(unique(CMP.begin(), CMP.end()), CMP.end());
	for(int i = 1; i <= n; ++i) C[i] = upper_bound(CMP.begin(), CMP.end(), C[i]) - CMP.begin();
	
	for(int i = 0; i < n - 1; ++i) {
		scanf("%d%d", A + i, B + i); 
		G[A[i]].PB(B[i]);
		P[B[i]] = A[i];
	}

	dfs(1); 
		
//	for(int i = 1; i <= n; ++i) printf("dep %d siz %d ch %d\n", DEP[i], SIZ[i], CH[i]);

	S[CH[1]].insert({DEP[1], C[1]});
	for(int i = 0; i < n - 1; ++i) {
		S[CH[B[i]]].insert({DEP[B[i]], C[B[i]]});
		printf("%lld\n", solve(A[i], C[B[i]]));
	}
	return 0;
}

Compilation message (stderr)

construction.cpp: In function 'int main()':
construction.cpp:99:7: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   99 |  scanf("%d", &n);
      |  ~~~~~^~~~~~~~~~
construction.cpp:101:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  101 |   scanf("%d", C + i);
      |   ~~~~~^~~~~~~~~~~~~
construction.cpp:110:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
  110 |   scanf("%d%d", A + i, B + i);
      |   ~~~~~^~~~~~~~~~~~~~~~~~~~~~

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