Submission #12132353


Source Code Expand

Copy
//Konrad Paluszek,University of Warsaw(former XIV LO Staszic)
//#STAY AT HOME
#ifndef LOCAL
#pragma GCC optimize("O3")
#endif
#define TIME (chrono::steady_clock::now().time_since_epoch().count())
#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#define xfm(a,b)a##b
#define xwb(a,b)xfm(a,b)
#define _ xwb(nvj,__LINE__)
#define __ xwb(kjy,__LINE__)
#define ___ xwb(cjm,__LINE__)
#define REP(i,n)for(urs(n)i=0;i<(n);++i)
#define UNTIL(t)while(clock()<(t)*CLOCKS_PER_SEC)
#define PER(r...)for(bool _=1;_||next_permutation(r);_=false)
#define ALL(r)(r).begin(),(r).end()
#define RALL(r)(r).rbegin(),(r).rend()
#define FS(r)r.first,r.second
#define SF(r)r.second,r.first
#define M0(r) memset(r, 0, sizeof(r))
#define sim template<class c
#define ros return o
#define rans return ans
#define forbits(i,m)if(m)for(urs(m)i=ctz(m),i##nvj=m;i##nvj;i##nvj^=((urs(m))1<<i),i=ctz(i##nvj))
#define fordbits(i,m)if(m)for(urs(m)i=8*sizeof(m)-clz(m)-1,i##nxd=m;i##nxd;i##nxd^=((urs(m))1<<i),i=8*sizeof(m)-clz(i##nxd)-1)
#define ksets(t, m, k, n) for (t m = (((t)1 << (k)) - 1); m < ((t)1 << (n)); m = nux(m))
#define urs(r...)typename decay<decltype(r)>::type
#define hur(f,g,r)sim>int f(c a){if(sizeof(c)==16)return r;if(sizeof(c)==8)return g##ll(a);return g(a);}
#define pwq(t,i) int clz(t x){return clz<int>(x)-i;}
using namespace __gnu_pbds;using namespace std;using ll=long long;using ld=long double;using ull=unsigned long long;using vi=vector<int>;using pii=pair<int,int>;using pll=pair<ll,ll>;using vpii=vector<pii>;using spii=set<pii>;using mii=map<int,int>;using unt=unsigned int;sim>using min_queue=priority_queue<c,vector<c>,greater<c>>;sim,class b,class cmp=less<c> >using ordered_map=tree<c,b,cmp,rb_tree_tag,tree_order_statistics_node_update>;sim, class cmp=less<c>>using ordered_set = ordered_map<c,null_type,cmp>;hur(popc,__builtin_popcount,popc<ull>(a)+popc<ull>(a>>64))hur(ctz,__builtin_ctz,(ull)a?ctz<ull>(a):64+ctz<ull>(a>>64))hur(clz,__builtin_clz,a>>64?clz<ull>(a>>64):64+clz<ull>(a))pwq(short,16)pwq(uint16_t,16)pwq(char,24)pwq(int8_t,24)pwq(uint8_t,24)sim,class N>bool mini(c&o,const N&h){if(o>h)ros=h,1;return 0;}sim,class N>bool maxi(c&o,const N&h){if(o<h)ros=h,1;return 0;}sim,class n>using gyv=c;
#if defined(LOCAL) // || defined(LOCAL2)
#include </home/kjp/headers/debuglib.hpp>
#else
#define loc(...)
#define onl(r...)r
#define debug(...)
#define print_stack(...)
#define mark_stack(...)
#define set_pre(...)
#define reg_it(...)
#define def_op(...) struct _{};
#if !defined(LOCAL) && !defined(LOCAL2)
#define exit my_exit
void my_exit(int x) {fflush(stdout);_Exit(x);}
#endif
#endif
#define next nexT
#define prev preV
#define tree trEE
#define left lefT
#define right righT
#define div diV
#define y1 y_1
ull mix(ull o){o+=0x9e3779b97f4a7c15;o=(o^(o>>30))*0xbf58476d1ce4e5b9;o=(o^(o>>27))*0x94d049bb133111eb;ros^(o>>31);}ull SALT=0x7a14a4b0881ebf9,tqu=0x7a14a4b0881ebf9;ull my_rand(){return tqu=mix(tqu);}void my_srand(ull x){SALT=tqu=x;}const int inf = 1023400000;const ll llinf=1234567890000000000ll;ll fix(ll o, ll m){o%=m;if(o<0)o+=m;ros;}
#define rand my_rand
#define srand my_srand
#define random_shuffle(r...)random_shuffle(r,[](int _){return my_rand()%_;})
sim>inline c nux(c m){if(!m)return numeric_limits<c>::max();c A=m&-m;c B=~((A-1)^m);c C=B&-B;c D=(C>>(1+ctz(A)))-1;return C|(m&~(C-1))|D;}__attribute__((no_sanitize_undefined))ll mul(ll a,ll b,ll m){ll q=a*(ld)b/m;ll o=a*b-q*m;o%=m;if(o<0)o+=m;ros;}sim>void unq(c&x){x.resize(unique(ALL(x))-x.begin());}
#if ((ULONG_MAX) != (UINT_MAX))
namespace std {template <> struct is_signed<__int128> : public true_type {};}
#endif
sim, class d> typename common_type<c,d>::type floor_div(c a, d b) {static_assert(is_signed<c>::value==is_signed<d>::value,"using floor_div with different signedness");if (b < 0) b = -b, a = -a;return a / b - (a % b < 0);}
sim, class d> typename common_type<c,d>::type ceil_div(c a, d b) {static_assert(is_signed<c>::value==is_signed<d>::value,"using ceil_div with different signedness");if (b < 0) b = -b, a = -a;return a / b + (a % b > 0);}
sim> struct REV {using value_type=typename c::value_type;c &x;using it=typename c::reverse_iterator;it begin(){return x.rbegin();}it end(){return x.rend();}};
sim> struct CREV {using value_type=typename c::value_type;const c&x;using it=typename c::const_reverse_iterator;it begin(){return x.rbegin();}it end(){return x.rend();}};
sim> REV<c> reversed(c&x) {return REV<c>{x};}sim> CREV<c> reversed(const c&x) {return CREV<c>{x};}
#define done(r...) exit(0 * printf(r))
#if defined(LOCAL) || defined(LOCAL2)
void __tmi() {cerr << setprecision(6) << fixed << "total time: " << clock() / (ld)CLOCKS_PER_SEC << "s" << endl;}int _ = (atexit(__tmi), 0);
#endif
//#STAY AT HOME
const int nax = 2e5 + 44;
vi graf[nax];
int col[nax];
int in[nax], out[nax];
int siz[nax];
int p[nax];
void dfs(int x, int y = 0) {
	p[x] = y;
	static int timer = 0;
	siz[x] = 1;
	in[x] = timer++;
	for (int v : graf[x]) if (v != y) {
		dfs(v, x);
		siz[x] += siz[v];
	}
	debug(arr(siz, x));
	out[x] = timer;
}
set <pii> by_col[nax];
int n;
ll solve(set <pii> &curr) {
	int left = n - curr.size();
	auto it = curr.end();
	ll ans = n * (n + 1ll) / 2;
	auto inc = [&](int x) {
		debug(imie(x));
		assert(x >= 0);
		left -= x;
		ans -= x * (x + 1ll) / 2;
		assert(left >= 0);
	};
	while (true) {
		if (it == curr.begin()) break;
		it--;
		auto x = it->second;
		for (int v : graf[x]) if (v != p[x]) {
			int curr_size = siz[v];
			int low = in[v], high = out[v];
			auto start = curr.lower_bound(pii(low, -1));
			reg_it(curr);
			debug(imie(start));
			auto end = start;
			while (end != curr.end() && end->first < high) {
				debug("decrease by", arr(siz, end->second));
				curr_size -= siz[end->second];
				end++;
			}
			debug(imie(x), imie(v), imie(curr_size), "from", range(start, end));
			curr.erase(start, end);
			inc(curr_size);
		}
	}
	inc(left);
	rans;
}
int main() {
	scanf("%d", &n);
	for (int i = 1; i <= n; ++i) scanf("%d", col + i);
	REP(_, n - 1) {
		int a, b;
		scanf("%d%d", &a, &b);
		graf[a].push_back(b);
		graf[b].push_back(a);
	}
	dfs(1);
	debug(draw_tree(1, graf, [](int x){stringstream o; o << "[" << x << ", " << in[x] << "; c=" << col[x] << "; s=" << siz[x] << "]"; ros.str();}));
	for (int i = 1; i <= n; ++i) by_col[col[i]].emplace(in[i], i);
	for (int i = 1; i <= n; ++i) {
		debug(arr(by_col, i));
		printf("%lld\n", solve(by_col[i]));
	}
	exit(0);
}
//#STAY AT HOME

Submission Info

Submission Time
Task F - path pass i
User i_love_chickpea
Language C++ (GCC 9.2.1)
Score 600
Code Size 6610 Byte
Status AC
Exec Time 283 ms
Memory 52696 KB

Compile Error

./Main.cpp: In function ‘int main()’:
./Main.cpp:130:7: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
  130 |  scanf("%d", &n);
      |  ~~~~~^~~~~~~~~~
./Main.cpp:131:36: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
  131 |  for (int i = 1; i <= n; ++i) scanf("%d", col + i);
      |                               ~~~~~^~~~~~~~~~~~~~~
./Main.cpp:134:8: warning: ignoring return value of ‘int scanf(const char*, ...)’, declared with attribute warn_unused_result [-Wunused-result]
  134 |   scanf("%d%d", &a, &b);
      |   ~~~~~^~~~~~~~~~~~~~~~

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 600 / 600
Status
AC × 5
AC × 24
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All line_01.txt, line_3_01.txt, random_01.txt, random_02.txt, random_03.txt, random_100_01.txt, random_100_02.txt, random_10_01.txt, random_10_02.txt, random_1_01.txt, random_1_02.txt, random_2_01.txt, random_2_02.txt, random_2_03.txt, random_3_01.txt, random_3_02.txt, random_3_03.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt, star_01.txt, star_3_01.txt
Case Name Status Exec Time Memory
line_01.txt AC 188 ms 52688 KB
line_3_01.txt AC 143 ms 52696 KB
random_01.txt AC 225 ms 37160 KB
random_02.txt AC 226 ms 37308 KB
random_03.txt AC 226 ms 37316 KB
random_100_01.txt AC 257 ms 37332 KB
random_100_02.txt AC 263 ms 37268 KB
random_10_01.txt AC 252 ms 37168 KB
random_10_02.txt AC 261 ms 37308 KB
random_1_01.txt AC 276 ms 37228 KB
random_1_02.txt AC 273 ms 37264 KB
random_2_01.txt AC 283 ms 37292 KB
random_2_02.txt AC 254 ms 37120 KB
random_2_03.txt AC 259 ms 37180 KB
random_3_01.txt AC 251 ms 37332 KB
random_3_02.txt AC 256 ms 37140 KB
random_3_03.txt AC 264 ms 37192 KB
sample_01.txt AC 15 ms 17612 KB
sample_02.txt AC 15 ms 17528 KB
sample_03.txt AC 13 ms 17608 KB
sample_04.txt AC 15 ms 17664 KB
sample_05.txt AC 18 ms 17656 KB
star_01.txt AC 149 ms 38024 KB
star_3_01.txt AC 118 ms 37824 KB