Submission #97538

#	Time	Username	Problem	Language	Result	Execution time	Memory
97538		IgorBaliuk	Transport (COCI19_transport)	C++14	130 / 130	365 ms	13996 KiB

Main

//#pragma GCC target("sse,sse2,sse3,ssse3,sse4,popcnt,abm,mmx,avx,tune=native")
//#pragma GCC optimize("unroll-loops")

#define _CRT_SECURE_NO_WARNINGS
#include <bits/stdc++.h>

using namespace std;

typedef pair <int, int> pii;
typedef pair <long long, long long> pll;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;

#define endl "\n"
#define mt make_tuple
#define mp make_pair

template <typename T1, typename T2>
bool umax(T1 &a, const T2&b) { return a < b ? a = b, 1 : 0;}
template <typename T1, typename T2> 
bool umin(T1 &a, const T2 &b) {return a > b ? a = b, 1 : 0;}
template <typename T>
T sqr(T a) {return a * a;}

mt19937 rng(20010709);
const int mod = 1000000007;
const int INF = 1000000001;
const ll INFLL = INF;
const int N = 100005;
const int M = 17;

struct Centroid {
	int n, t;
	vector <vector <pii>> e;
	vector <int> locked, a, len;
	vector <ll> globalUp, globalDown, localUp, localDown;
	ll result;

	Centroid() { }
	void init(int N) {
		n = N;
		e.assign(n + 1, vector <pii>());
		locked.assign(n + 1, 0);
		a.assign(n + 1, 0);
		len.assign(n + 1, 0);
		result = 0;
	}

	int dfsLen(int v, int p = -1) {
		len[v] = 1;

		for (auto it : e[v]) {
			if (!locked[it.first] && (it.first != p))
				len[v] += dfsLen(it.first, v);
		}

		return len[v];
	}

	int findCentroid(int v) {
		int p = -1, found = 1, fullLen = len[v];

		while (found) {
			found = 0;

			for (auto it : e[v]) {
				if (!locked[it.first] && (it.first != p) && (len[it.first] * 2 >= fullLen)) {
					p = v;
					v = it.first;
					found = 1;
					break;
				}
			}
		}

		return v;
	}

	void process(vector <ll> &up, vector <ll> &down, int d = 1) {
		sort(up.begin(), up.end());
		sort(down.begin(), down.end());

		for (int i = 0, j = 0; i < down.size(); i++) {
			while ((j < up.size()) && (up[j] < down[i]))
				j++;

			result += d * (1LL * up.size() - j);
		}
	}

	void dfs(int v, int p, ll up, ll maxUp, ll down, ll maxDown) {
		umax(maxDown, -down);
		up += a[v];
		down += a[v];

		if (a[v] >= maxUp)
			localUp.push_back(up);

		localDown.push_back(maxDown);

		for (auto it : e[v]) {
			if (!locked[it.first] && (it.first != p))
				dfs(it.first, v, up - it.second, it.second + max(0LL, maxUp - a[v]), down - it.second, maxDown);
		}
	}

	void get(int v) {
		dfsLen(v);
		v = findCentroid(v);
		t = a[v];

		globalUp.clear();
		globalDown.clear();
		globalUp.push_back(0);
		globalDown.push_back(-a[v]);
		result--;

		for (auto it : e[v]) {
			if (!locked[it.first]) {
				localUp.clear();
				localDown.clear();

				dfs(it.first, v, -it.second, it.second, (a[v] - it.second), 0);

				process(localUp, localDown, -1);
				globalUp.insert(globalUp.end(), localUp.begin(), localUp.end());
				globalDown.insert(globalDown.end(), localDown.begin(), localDown.end());
			}
		}

		process(globalUp, globalDown);

		locked[v] = 1;

		for (auto it : e[v]) {
			if (!locked[it.first])
				get(it.first);
		}
	}

	ll solve() {
		get(1);
		return result;
	}
};

int n;
Centroid c;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

	cin >> n;

	c.init(n);

	for (int i = 1; i <= n; i++)
		cin >> c.a[i];

	for (int i = 1; i < n; i++) {
		int u, v, w;
		cin >> u >> v >> w;
		c.e[u].push_back({ v, w });
		c.e[v].push_back({ u, w });
	}

    cout << c.solve() << endl;

    return 0;
}

Compilation message (stderr)

transport.cpp: In member function 'void Centroid::process(std::vector<long long int>&, std::vector<long long int>&, int)':
transport.cpp:84:28: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for (int i = 0, j = 0; i < down.size(); i++) {
                          ~~^~~~~~~~~~~~~
transport.cpp:85:14: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
    while ((j < up.size()) && (up[j] < down[i]))
            ~~^~~~~~~~~~~

• Subtask 1: 13 / 13
• Subtask 2: 13 / 13
• Subtask 3: 13 / 13
• Subtask 4: 13 / 13
• Subtask 5: 13 / 13
• Subtask 6: 13 / 13
• Subtask 7: 13 / 13
• Subtask 8: 13 / 13
• Subtask 9: 13 / 13
• Subtask 10: 13 / 13