Submission #1243586

#	Time	Username	Problem	Language	Result	Execution time	Memory
1243586		Sam_arvandi	Capital City (JOI20_capital_city)	C++20	100 / 100	465 ms	165832 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;

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

#define FOR(i, j, n) for (int i = j; i<= n; i++)
#define ROF(i, n, j) for (int i = n; i>= j; i--)
#define pb push_back
#define mp make_pair
#define S second
#define F first
#define IOS ios_base :: sync_with_stdio(0); cin.tie(0); cout.tie(0)

/*
#pragma GCC optimization("Ofast, unroll-loops")
#pragma GCC target("avx2")
#pragma GCC target("bmi")
#pragma GCC target("bmi2")
#pragma GCC target("lzcnt")
*/

const int mn = 2e5 + 5, mlg = 22, mn2 = 6e5 + 10;
vector<int> A[mn2], a[mn], A2[mn2];
vector<pii> a2[mn];
int h[mn], siz[mn], par[mn], nemd[mn], beg[mn], c[mn], lift[mn][mlg];
int root[mn], g[mn], Siz[mn];
int k[mn], num[mn2];
bool mark[mn], MARK[mn2], MARK2[mn2], flag[mn2];
vector<int> melat[mn2];
struct Seg_tree
{
        struct Node
        {
                int lc = -1, rc = -1;
        }tmp;
        int n = 0;
        vector<Node> node;

        void init(int id, int L, int R)
        {
                if (L+1 == R)
                {
                        A[id].pb(nemd[L]);
                        return;
                }
                n++;
                node.pb(tmp);
                node[id].lc = n;
                A[id].pb(n);
                n++;
                node.pb(tmp);
                node[id].rc = n;
                A[id].pb(n);
                int mid = (L+R)/2;
                init(node[id].lc, L, mid);
                init(node[id].rc, mid, R);
        }


        void add(int id, int L, int R, int l, int r, int u)
        {
                if (L == l and R == r)
                {
                        A[u].pb(id);
                        return;
                }
                int mid = (L+R)/2;
                if (l >= mid) add(node[id].rc, mid, R,l, r, u);
                else if (r <= mid) add(node[id].lc, L, mid, l, r, u);
                else
                {
                        add(node[id].lc, L, mid, l, mid, u);
                        add(node[id].rc, mid, R, mid, r, u);
                }
        }
}seg;

void dfs(int u)
{
        mark[u] =1;
        siz[u] =1;
        for(auto v: a[u])
        {
                if (mark[v])
                {
                        par[u] = v;
                        lift[u][0] = v;
                        continue;
                }
                h[v] = h[u]+1;
                dfs(v);
                siz[u] += siz[v];
                a2[u].pb(mp(siz[v], v));
        }
}

void dfs2(int u)
{
        int v;
        for(auto p: a2[u])
        {
                v = p.S;
                if (v == a2[u][0].S)
                {
                        root[v] = root[u];
                }
                else root[v] = v;
                dfs2(v);
        }
}

vector<int> alaki;
void DFS(int u)
{
        MARK[u] = 1;
        for(auto v: A[u])
        {
                if (!MARK[v]) DFS(v);
        }
        alaki.pb(u);
}

int ind = 0;
void DFS2(int u)
{
        num[u] = ind;
        melat[ind].pb(u);
        MARK2[u] = 1;
        for(auto v: A2[u])
        {
                if (!MARK2[v]) DFS2(v);
        }
        return;
}
struct HLD
{
        void init(int n)
        {
                FOR(i, 1, n)
                {
                        if (a2[i].size() != 0) continue;
                        Siz[root[i]] = 0;
                        int u = i;
                        while (root[u] != u)
                        {
                                Siz[root[u]]++;
                                g[u] = Siz[root[u]];
                                nemd[Siz[root[u]]] = c[u];
                                u = par[u];
                        }
                        Siz[root[u]]++;
                        g[u] = Siz[root[u]];
                        nemd[Siz[root[u]]] = c[u];
                        seg.n++;
                        beg[u] = seg.n;
                        seg.node.pb(seg.tmp);
                        seg.init(seg.n, 1, Siz[u]+1);
                }
        }

        void update(int u, int v, int x)
        {
                if (h[root[u]] <= h[v])
                {
                        seg.add(beg[root[u]], 1, Siz[root[u]]+1, g[u], g[v]+1, x);
                }
                else
                {
                        seg.add(beg[root[u]], 1, Siz[root[u]]+1, g[u], Siz[root[u]]+1, x);
                        update(par[root[u]], v, x);
                }
                return;
        }
}hld;

int get_lca(int u, int v)
{
        if (h[u] < h[v]) swap(u, v);
        ROF(i, 20, 0)
        {
                if (h[lift[u][i]] >= h[v])
                {
                        u = lift[u][i];
                }
        }
        if (u == v) return u;
        ROF(i, 20, 0)
        {
                if (lift[u][i] != lift[v][i])
                {
                        u = lift[u][i];
                        v = lift[v][i];
                }
        }
        return par[u];
}


signed main()
{
        IOS;
        int n, m, u, v;
        cin >> n >> m;
        FOR(i,1 , n-1)
        {
                cin >> u >> v;
                a[u].pb(v);
                a[v].pb(u);
        }
        FOR(i, 1, n)
        {
                cin >>c[i];
        }
        FOR(i, 0, m)
        {
                seg.n++;
                seg.node.pb(seg.tmp);
        }
        h[1] = 1;
        dfs(1);
        FOR(j, 1, 20)
        {
                FOR(i, 1, n)
                {
                        lift[i][j] = lift[lift[i][j-1]][j-1];
                }
        }
        FOR(i,1 , n)
        {
                sort(a2[i].begin(), a2[i].end());
                reverse(a2[i].begin(), a2[i].end());
        }
        root[1] = 1;
        dfs2(1);
        //return 0;
        FOR(i, 1, n)
        {
                if (k[c[i]] == 0) k[c[i]] = i;
                else k[c[i]] = get_lca(i, k[c[i]]);
        }
        
        hld.init(n);
        FOR(i, 1, n)
        {
                hld.update(i, k[c[i]], c[i]);
        }
        //return 0;
        int N = seg.n;

        FOR(i, 1, N)
        {
                for(auto v: A[i])
                {
                        A2[v].pb(i);
                }
        }
        FOR(i, 1, N)
        {
                if (!MARK[i])
                {
                        //cout << i << endl;
                        DFS(i);
                }
        }
        reverse(alaki.begin(), alaki.end());
        for(auto i: alaki)
        {
                if (!MARK2[i])
                {
                        ind++;
                        DFS2(i);
                }
        }
        FOR(i, 1, N)
        {
                for(auto v: A[i])
                {
                        if (num[v] != num[i]) flag[num[i]] = 1;
                }
        }
        int ans = N, ted =0;
        FOR(i, 1, ind)
        {
                if (flag[i]) continue;
                ted = 0;
                for(auto v: melat[i])
                {
                        if (v > m) continue;
                        ted++;
                }
                if (ted != 0) ans = min(ans, ted);
        }
        cout << ans-1;
        return 0;
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59