Submission #655785

#	Time	Username	Problem	Language	Result	Execution time	Memory
655785		aZvezda	Mergers (JOI19_mergers)	C++14	100 / 100	770 ms	72300 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
//#pragma GCC optimize ("O3")
//#pragma GCC target ("sse4")
#define endl "\n"
typedef long long ll;
template<class T, class T2> inline ostream &operator <<(ostream &out, const pair<T, T2> &x) { out << x.first << " " << x.second; return out;}
template<class T, class T2> inline istream &operator >>(istream &in, pair<T, T2> &x) { in >> x.first >> x.second; return in;}
template<class T, class T2> inline bool chkmax(T &x, const T2 &y) { return x < y ? x = y, 1 : 0; }
template<class T, class T2> inline bool chkmin(T &x, const T2 &y) { return x > y ? x = y, 1 : 0; }
const ll mod = 1e9 + 7;
#define out(x) "{" << (#x) << ": " << x << "} "

const int MAX_N = 5e5 + 10;
void *arenaPtr;
size_t arenaSize;
template<class T>
struct MyAllocator {
    typedef T value_type;

    MyAllocator () = default;
    template<class U>
    constexpr MyAllocator (const MyAllocator <U>&) noexcept {}

	T* allocate(size_t n) {
		const auto ret = (T*)((char*)arenaPtr + arenaSize);
		arenaSize += n * sizeof(T);
		return ret;
    }
	void deallocate(T* p, size_t n) {
    	return;
	}
};

vector<int , MyAllocator<int> > g[MAX_N], nw[MAX_N];
int cnt[MAX_N], cntNow[MAX_N], col[MAX_N];
int n, k, ret;
int sz[MAX_N], big[MAX_N];
int cntok;

void count(int x, int delta = 1) {
    if(cntNow[col[x]] == 0 || cntNow[col[x]] == cnt[col[x]]) {
        cntok --;
    }
    cntNow[col[x]] += delta;
    if(cntNow[col[x]] == 0 || cntNow[col[x]] == cnt[col[x]]) {
        cntok ++;
    }
}

void add(int x, int p, int delta = 1) {
    count(x, delta);
    for(auto it : g[x]) {
        if(it == p) {continue;}
        add(it, x, delta);
    }
}

void calcSize(int x, int p) {
    sz[x] = 1;
    big[x] = 0;
    for(auto it : g[x]) {
        if(it == p) {continue;}
        calcSize(it, x);
        sz[x] += sz[it];
        if(sz[it] > sz[big[x]]) {
            big[x] = it;
        }
    }
}

vector<int, MyAllocator<int> > stck;
void dfs(int x, int p, bool keep = true) {
    int topRet = stck.size();

    for(auto it : g[x]) {
        if(it == p || it == big[x]) {continue;}
        dfs(it, x, false);
    }
    if(big[x] != 0) {
        dfs(big[x], x, true);
    }

    for(auto it : g[x]) {
        if(it == p || it == big[x]) {continue;}
        add(it, x, 1);
    }
    count(x);

    if(cntok == k) {
        while(stck.size() > topRet) {
            nw[x].push_back(stck.back());
            stck.pop_back();
        }
        stck.push_back(x);
    }

    if(!keep) {
        add(x, p, -1);
    }
}

void calc(int x, int d) {
    if(x == 1 && nw[x].size() == 1) {
        ret ++;
    } else if(nw[x].size() == 0) {
        ret ++;
    }
    for(auto it : nw[x]) {
        calc(it, d + 1);
    }
}

signed main() {
	arenaPtr = malloc(MAX_N * 1000);
	arenaSize = 0;

	ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
    cin >> n >> k;
    cntok = k;
    for(int i = 0; i < n - 1; i ++) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    for(int i = 1; i <= n; i ++) {
        cin >> col[i];
        cnt[col[i]] ++;
    }
    calcSize(1, 0);
    dfs(1, 0, 1);
    calc(1, 0);
    if(ret != 1) {
        cout << (ret + 1) / 2 << endl;
    } else {
        cout << 0 << endl;
    }

	free(arenaPtr);
    return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'void dfs(int, int, bool)':
mergers.cpp:91:27: warning: comparison of integer expressions of different signedness: 'std::vector<int, MyAllocator<int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   91 |         while(stck.size() > topRet) {
      |               ~~~~~~~~~~~~^~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30