Submission #153059

#	Time	Username	Problem	Language	Result	Execution time	Memory
153059		luciocf	Transport (COCI19_transport)	C++14	130 / 130	540 ms	15932 KiB

Main

#include <bits/stdc++.h>

using namespace std;

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

const int maxn = 1e5+10;

int n;
ll ans;

int a[maxn];

int bit[maxn], id[maxn];

int sz[maxn];

bool mark[maxn];

vector<pii> grafo[maxn]; 
vector<pair<ll, int>> tot;

void upd(int x, int v)
{
	for (; x <= n; x += (x&-x))
		bit[x] += v;
}

int soma(int x)
{
	int ans = 0;
	for (; x > 0; x -= (x&-x))
		ans += bit[x];
	return ans;
}

void dfs(int u, int p)
{
	sz[u] = 1;

	for (auto pp: grafo[u])
	{
		int v = pp.first, w = pp.second;
		if (v == p || mark[v]) continue;

		dfs(v, u);

		sz[u] += sz[v];
	}
}

int centroid(int u, int p, int S)
{
	for (auto pp: grafo[u])
	{
		int v = pp.first;
		if (v == p || mark[v]) continue;

		if (sz[v] > S/2) return centroid(v, u, S);
	}

	return u;
}

void dfs_up(int u, int p, ll cost, ll sum)
{
	if (cost >= 0)
	{
		ans++;
		tot.push_back({sum, u});
	}

	for (auto pp: grafo[u])
	{
		int v = pp.first, w = pp.second;
		if (v == p || mark[v]) continue;

		dfs_up(v, u, min(cost+1ll*a[v]-1ll*w, 1ll*a[v]-1ll*w), sum+1ll*a[v]-1ll*w);
	}
}

void dfs_add(int u, int p, ll cost, int add)
{
	if (cost >= 0)
		upd(id[u], add);

	for (auto pp: grafo[u])
	{
		int v = pp.first, w = pp.second;
		if (v == p || mark[v]) continue;

		dfs_add(v, u, min(cost+1ll*a[v]-1ll*w, 1ll*a[v]-1ll*w), add);
	}
}

void dfs_down(int u, int p, ll cost, ll sum)
{
	if (cost >= 0)
		ans++;

	if (lower_bound(tot.begin(), tot.end(), make_pair(-cost, 0)) != tot.end())
	{
		int pos = (int) (lower_bound(tot.begin(), tot.end(), make_pair(-cost, 0))-tot.begin())+1;
		ans += 1ll*(soma(n)-soma(pos-1));
	}

	for (auto pp: grafo[u])
	{
		int v = pp.first, w = pp.second;
		if (v == p || mark[v]) continue;

		dfs_down(v, u, min(cost, sum+1ll*a[u]-1ll*w), sum+1ll*a[u]-1ll*w);
	}
}

void decompose(int u)
{
	dfs(u, 0);

	int c = centroid(u, 0, sz[u]);
	mark[c] = 1;

	for (auto v: grafo[c])
		if (!mark[v.first])
			dfs_up(v.first, c, 1ll*(a[v.first]-v.second), 1ll*(a[v.first]-v.second));

	sort(tot.begin(), tot.end());

	for (int i = 0; i < tot.size(); i++)
		id[tot[i].second] = i+1;

	for (auto v: grafo[c])
		if (!mark[v.first])
			dfs_add(v.first, c, 1ll*(a[v.first]-v.second), 1);

	for (auto pp: grafo[c])
	{
		int v = pp.first, w = pp.second;
		if (mark[v]) continue;

		dfs_add(v, c, 1ll*(a[v]-w), -1);

		dfs_down(v, c, 1ll*(a[c]-w), 1ll*(a[c]-w));

		dfs_add(v, c, 1ll*(a[v]-w), 1);
	}

	for (auto v: grafo[c])
		if (!mark[v.first])
			dfs_add(v.first, c, 1ll*(a[v.first]-v.second), -1);

	tot.clear();

	for (auto v: grafo[c])
		if (!mark[v.first])
			decompose(v.first);
}

int main(void)
{
	scanf("%d", &n);

	for (int i = 1; i <= n; i++)
		scanf("%d", &a[i]);

	for (int i = 1; i < n; i++)
	{
		int u, v, w;
		scanf("%d %d %d", &u, &v, &w);

		grafo[u].push_back({v, w});
		grafo[v].push_back({u, w});
	}

	decompose(1);

	printf("%lld\n", ans);
}

Compilation message (stderr)

transport.cpp: In function 'void dfs(int, int)':
transport.cpp:44:21: warning: unused variable 'w' [-Wunused-variable]
   int v = pp.first, w = pp.second;
                     ^
transport.cpp: In function 'void decompose(int)':
transport.cpp:130:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = 0; i < tot.size(); i++)
                  ~~^~~~~~~~~~~~
transport.cpp: In function 'int main()':
transport.cpp:162:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d", &n);
  ~~~~~^~~~~~~~~~
transport.cpp:165:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d", &a[i]);
   ~~~~~^~~~~~~~~~~~~
transport.cpp:170:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d %d", &u, &v, &w);
   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~

• 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