Submission #293521

#	Time	Username	Problem	Language	Result	Execution time	Memory
293521		Aldas25	Mergers (JOI19_mergers)	C++14	100 / 100	1768 ms	165392 KiB

mergers

//#pragma GCC optimize("O3")

#include <bits/stdc++.h>
//#include <ctime>

using namespace std;

#define FAST_IO ios_base::sync_with_stdio(0); cin.tie(nullptr)
#define FOR(i, a, b) for(int i = (a); i <= (b); i++)
#define REP(n) FOR(O, 1, (n))
#define pb push_back
#define f first
#define s second
typedef long double ld;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<int, pii> piii;
typedef vector<int> vi;
typedef vector<pii> vii;
typedef vector<ll> vl;
typedef vector<piii> viii;
//mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MAXN = 500100, MAXK = 20;
const ll MOD = 1e9+7;
const ll INF = 1e16;
const ld PI = asin(1) * 2;

void setIO () {
    FAST_IO;
}

void setIO (string s) {
    setIO();
 	freopen((s+".in").c_str(),"r",stdin);
 	freopen((s+".out").c_str(),"w",stdout);
}

int n, k;
vi adj[MAXN];
int s[MAXN];
vi bel[MAXN];

int deg[MAXN];

int par[MAXN];
int find (int a) {return par[a] = par[a]==a ? a : find(par[a]);}
bool same (int a, int b) {return find(a) == find(b);}
void unite (int a, int b) {
    a = find(a), b = find(b);
    if (a == b) return;
    par[b] = a;
}

int curT = 0;
vi times[MAXN];
int pr[MAXN][MAXK];
int dep[MAXN];

void dfs1 (int v, int p) {
    times[v].pb(++curT);

    FOR(i, 1,MAXK-1)
        pr[v][i] = pr[pr[v][i-1]][i-1];

    for (int u : adj[v]) {
        if (u == p) continue;
        pr[u][0] = v;
        dep[u] = dep[v]+1;
        dfs1 (u, v);
        times[v].pb(++curT);
    }
}

int lca (int v, int u) {
    if (u == v) return u;
    if (dep[u] > dep[v]) swap(u,v);
    int d = dep[v] - dep[u];
    FOR(i, 0, MAXK-1)
        if (d & (1<<i))
            v = pr[v][i];
    if (u == v) return u;
    for (int i = MAXK-1; i >= 0; i--) {
        if (pr[v][i] != pr[u][i]) {
            v = pr[v][i];
            u = pr[u][i];
        }
    }
    return pr[v][0];
}

int cnt[MAXN];

void dfs3 (int v, int p) {
    for (int u : adj[v]) {
        if (u == p) continue;
        dfs3(u,v);
        cnt[v] += cnt[u];
    }
//    cout << "  v = " << v << " cnt = " << cnt[v] << endl;

    if (p != -1 && cnt[v] > 0) unite(v,p);
}

void dfs2 (int v, int p) {
    for (int u : adj[v]) {
        if (u == p) continue;
        dfs2(u, v);
    }
    if (p == -1) return;
    if (!same(v,p)) {
        deg[find(v)]++;
        deg[find(p)]++;
    }
}

vii seq;

int main() {
    setIO();

   // ifstream cin ("in.txt");

    cin >> n >> k;
    REP(n-1) {
        int a, b; cin >> a >> b;
        adj[a].pb(b);
        adj[b].pb(a);
    }
    FOR(i, 1, n) cin >> s[i];
    FOR(i,1 , n) bel[s[i]].pb(i);

    FOR(i, 1, n) par[i] = i;
    dfs1 (1, -1);


    FOR(i, 1, k) {
        seq.clear();
        for (int v : bel[i])
            for (int t : times[v])
                seq.pb({t, v});
        sort(seq.begin(), seq.end());
        FOR(j, 0, (int)seq.size()-2) {
            int u = seq[j].s;
            int v = seq[j+1].s;
            cnt[u]++;
            cnt[v]++;
            cnt[lca(u,v)] -= 2;
        }
    }

    dfs3 (1, -1);

    dfs2 (1, -1);

    //FOR(i, 1, n) cout <<"   i = " << i << " deg = " << deg[i] << endl;
    //FOR(i, 1, n) cout << "  i = " << i << " par = " << par[i] << endl;

    int leaves = 0;
    FOR(i, 1, n) if (deg[i] == 1) leaves++;
    int ans = (leaves+1)/2;
    cout << ans <<"\n";

	return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'void setIO(std::string)':
mergers.cpp:35:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
   35 |   freopen((s+".in").c_str(),"r",stdin);
      |   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~
mergers.cpp:36:10: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
   36 |   freopen((s+".out").c_str(),"w",stdout);
      |   ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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