Submission #1222788

#	Time	Username	Problem	Language	Result	Execution time	Memory
1222788		The_Samurai	Capital City (JOI20_capital_city)	C++20	100 / 100	275 ms	93384 KiB

capital_city

// I stand with PALESTINE



//#pragma GCC optimize("Ofast")
//#pragma GCC optimize ("unroll-loops")
//#pragma GCC target("avx,avx2")

#include <bits/stdc++.h>
//#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
//using namespace __gnu_pbds;

// #define int long long
#define ff first
#define ss second
#define sz(a) (int) (a).size()
//template<class T, class C = null_type> using ordered_tree = tree<T, C, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long double ld;

namespace io {

    template<typename F, typename S>
    ostream &operator<<(ostream &os, const pair<F, S> &p) { return os << p.ff << " " << p.ss; }

    template<typename F, typename S>
    ostream &operator<<(ostream &os, const map<F, S> &mp) {
        for (auto it: mp) { os << it << endl; }
        return os;
    }

    template<typename F>
    ostream &operator<<(ostream &os, const vector<F> &v) {
        bool space = false;
        for (F x: v) {
            if (space) os << " ";
            space = true;
            os << x;
        }
        return os;
    }

    template<typename F>
    ostream &operator<<(ostream &os, const deque<F> &d) {
        bool space = false;
        for (F x: d) {
            if (space) os << " ";
            space = true;
            os << x;
        }
        return os;
    }

    template<typename F>
    ostream &operator<<(ostream &os, const set<F> &st) {
        bool space = false;
        for (F x: st) {
            if (space) os << " ";
            space = true;
            os << x;
        }
        return os;
    }

    template<typename F>
    ostream &operator<<(ostream &os, const multiset<F> &st) {
        bool space = false;
        for (F x: st) {
            if (space) os << " ";
            space = true;
            os << x << x;
        }
        return os;
    }

    template<typename F, typename S>
    istream &operator>>(istream &is, pair<F, S> &p) { return is >> p.ff >> p.ss; }

    template<typename F>
    istream &operator>>(istream &is, vector<F> &v) {
        for (F &x: v) { is >> x; }
        return is;
    }

    long long fastread() {
        char c;
        long long d = 1, x = 0;
        do c = getchar(); while (c == ' ' || c == '\n');
        if (c == '-') c = getchar(), d = -1;
        while (isdigit(c)) {
            x = x * 10 + c - '0';
            c = getchar();
        }
        return d * x;
    }

    static bool sep = false;

    using std::to_string;

    string to_string(bool x) {
        return (x ? "true" : "false");
    }

    string to_string(const string &s) { return "\"" + s + "\""; }

    string to_string(const char *s) { return "\"" + string(s) + "\""; }

    string to_string(const char &c) {
        string s;
        s += c;
        return "\'" + s + "\'";
    }

    template<typename Type>
    string to_string(vector<Type>);

    template<typename F, typename S>
    string to_string(pair<F, S>);

    template<typename Collection>
    string to_string(Collection);

    template<typename F, typename S>
    string to_string(pair<F, S> p) { return "{" + to_string(p.ff) + ", " + to_string(p.ss) + "}"; }

    template<typename Type>
    string to_string(vector<Type> v) {
        bool sep = false;
        string s = "[";
        for (Type x: v) {
            if (sep) s += ", ";
            sep = true;
            s += to_string(x);
        }
        s += "]";
        return s;
    }

    template<typename Collection>
    string to_string(Collection collection) {
        bool sep = false;
        string s = "{";
        for (auto x: collection) {
            if (sep) s += ", ";
            sep = true;
            s += to_string(x);
        }
        s += "}";
        return s;
    }

    void print() {
        cout << endl;
        sep = false;
    }

    template<typename F, typename... Other>
    void print(F ff, Other... other) {
        if (sep) cout << " | ";
        sep = true;
        cout << to_string(ff);
        print(other...);
    }

}
using namespace io;

namespace utils {

    template<typename F, typename S>
    inline void maxs(F &a, S b) { a = a > b ? a : b; }

    template<typename F, typename S>
    inline void mins(F &a, S b) { a = a < b ? a : b; }

    template<typename F, typename S>
    long long max(F a, S b) { return a > b ? a : b; }

    template<typename F, typename S>
    long long min(F a, S b) { return a < b ? a : b; }

    constexpr long long operator "" _E(unsigned long long n) {
        long long p = 1, a = 10;
        for (int i = 0; i < n; i++) p *= a;
        return p;
    }

    random_device rd;
    mt19937 mt(rd());

    template<typename T>
    T rand(T l, T r) {
        uniform_int_distribution<T> dist(l, r);
        return dist(mt);
    };

}
using namespace utils;


#ifdef sunnatov
#define print(...) cout << "[" << #__VA_ARGS__ << "]: "; io::print( __VA_ARGS__ );
#else
#define print( ... ) 42
#endif

const int mod = 9_E + 7;
const double EPS = 1e-7;
long long doxuya = 18_E + 10;
int INF = 9_E + 10;
const char nl = '\n';

int month[] = {0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31};

/*

*/

template<typename T> struct Fenwick {
    int n;
    vector<T> tree;

    void init(int len) {
        n = len;
        tree.assign(n + 1, 0);
    }

    void update(int i, T v) {
        for (i++; i <= n; i += i & (-i)) tree[i] += v;
    }

    T get(int i) {
        T sum = 0;
        for (i++; i > 0; i -= i & (-i)) sum += tree[i];
        return sum;
    }

    T get(int l, int r) {
        return get(r) - get(l - 1);
    }
};  

struct hld {
    template<typename T> struct SegTree {
        vector<T> tree;
        int size;
        T neutral_element = 1e9; // sum - 0, mx - (-INF), mn - INF

        inline T merge(T a, T b) {
            return min(a, b);
        }

        void build(vector<T> &a) {
            size = 1;
            while (size < a.size()) size *= 2;
            tree.assign(2 * size, neutral_element);
            for (int i = size; i < size + a.size(); i++) tree[i] = a[i - size];
            for (int i = size - 1; i > 0; i--) tree[i] = merge(tree[i << 1], tree[i << 1 | 1]);
        }

        void set(int p, int value) {  // set value at position p
            p += size;
            tree[p] = merge(value, tree[p]);
            for (; p > 1; p >>= 1) tree[p >> 1] = merge(tree[p], tree[p ^ 1]);
        }

        T get(int l, int r) {  // sum on interval [l, r]
            if (l > r) return neutral_element;
            T res = neutral_element;
            for (l += size, r += size + 1; l < r; l >>= 1, r >>= 1) {
                if (l & 1) res = merge(res, tree[l++]);
                if (r & 1) res = merge(res, tree[--r]);
            }
            return res;
        }
    };

    vector<int> parent, depth, heavy, head, pos;
    SegTree<int> sg;
    int z;


    int dfs(int u, vector<vector<int>> &adj) {
        int size = 1, max_child_size = 0;
        for (int v: adj[u]) {
            if (parent[u] == v) continue;
            parent[v] = u; depth[v] = depth[u] + 1;
            int child_size = dfs(v, adj);
            size += child_size;
            if (max_child_size < child_size) {
                heavy[u] = v;
                max_child_size = child_size;
            }
        }
        return size;
    }

    void decompose(int u, int h, vector<vector<int>> &adj) {
        head[u] = h; pos[u] = z++;
        if (heavy[u] != -1) decompose(heavy[u], h, adj);
        for (int v: adj[u]) {
            if (v != parent[u] and v != heavy[u]) {
                decompose(v, v, adj);
            }
        }
    }

    void init(vector<vector<int>> &adj, vector<int> &a, int root) {
        int n = (int) adj.size();
        parent.assign(n, 0); depth.assign(n + 1, 0);
        heavy.assign(n + 1, -1); head.assign(n + 1, 0);
        pos.assign(n + 1, -1); z = 0;
        dfs(root, adj);
        decompose(root, root, adj);
        vector<int> vec(n);
        for (int i = 0; i < n; i++) {
            if (pos[i] != -1) vec[pos[i]] = a[i];
        }
        sg.build(vec);
    }

    int query(int a, int b) {
        int res = INF;
        for (; head[a] != head[b]; b = parent[head[b]]) {
            if (depth[head[a]] > depth[head[b]]) swap(a, b);
            res = min(res, sg.get(pos[head[b]], pos[b]));
        }
        if (depth[a] > depth[b]) swap(a, b);
        res = min(res, sg.get(pos[a], pos[b]));
        return res;
    }

    void upd(int a, int val) { sg.set(pos[a], val); }
};

void solution(istream &cin, ostream &cout, const int &test_case) {
    int n, m;
    cin >> n >> m;
    vector<vector<int>> g(n + 1);
    vector<int> color(n + 1);
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        g[u].emplace_back(v);
        g[v].emplace_back(u);
    }
    for (int i = 1; i <= n; i++) cin >> color[i];

    const int lg = 18;
    vector jump(lg, vector(n + 1, 0));
    vector<int> tin(n + 1), tout(n + 1), depth(n + 1), ord(n + 1);
    int z = 0;
    auto dfs = [&](auto &dfs, int u) -> void {
        for (int i = 1; i < lg; i++) {
            jump[i][u] = jump[i - 1][jump[i - 1][u]];
            if (!jump[i][u]) break;
        }
        tin[u] = tout[u] = z++;
        ord[tin[u]] = u;
        for (int v: g[u]) {
            if (jump[0][u] == v) continue;
            jump[0][v] = u;
            depth[v] = depth[u] + 1;
            dfs(dfs, v);
            tout[u] = tout[v];
        }
    };
    auto isPar = [&](int u, int v) { return tin[u] <= tin[v] and tout[v] <= tout[u]; };
    auto lca = [&](int u, int v) -> int {
        if (isPar(u, v)) return u;
        if (isPar(v, u)) return v;
        for (int i = lg - 1; i >= 0; i--) {
            if (jump[i][u] and !isPar(jump[i][u], v)) u = jump[i][u];
        }
        return jump[0][u];
    };
    dfs(dfs, 1);

    vector<vector<int>> vec(m + 1);
    vector<set<int>> st(m + 1);
    vector<int> need(n + 1);
    vector<bool> done(n + 1), dont(n + 1);
    Fenwick<int> fw; fw.init(n + 1);
    for (int i = 1; i <= n; i++) vec[color[i]].emplace_back(i);
    for (int i = 1; i <= m; i++) {
        sort(vec[i].begin(), vec[i].end(), [&](int u, int v) {
            return tin[u] < tin[v];
        });
        int highest = INF;
        for (int j = 0; j < sz(vec[i]); j++) mins(highest, depth[vec[i][j]]);
        for (int j = 1; j < sz(vec[i]); j++) mins(highest, depth[lca(vec[i][j - 1], vec[i][j])]);
        for (int j = 0; j < sz(vec[i]); j++) need[vec[i][j]] = highest;
        for (int j = 0; j < sz(vec[i]); j++) st[i].insert(tin[vec[i][j]]);
    }
    hld hld; hld.init(g, need, 1);

    auto dfs2 = [&](auto &dfs2, int u) -> void {
        for (int v: g[u]) {
            if (jump[0][u] == v) continue;
            dfs2(dfs2, v);
        }
        while (true) {
            auto it = st[color[u]].upper_bound(tin[u]);
            if (it == st[color[u]].end() or !isPar(u, ord[*it])) break;
            int v = ord[*it];
            mins(need[u], hld.query(u, v));
            if (fw.get(tin[v]) - fw.get(tin[u])) dont[u] = true;
            st[color[u]].erase(it);
        }
        hld.upd(u, need[u]);
        if (dont[u]) {
            fw.update(tin[u], 1);
            fw.update(tout[u] + 1, -1);
            return;
        }
        if (need[u] == depth[u]) {
            done[u] = true;
            fw.update(tin[u], 1);
            fw.update(tout[u] + 1, -1);
        }
    };
    dfs2(dfs2, 1);

    int ans = INF, tot = 0;
    for (int i = 1; i <= n; i++) {
        if (!done[i]) continue;
        vector<bool> vis(n + 1);
        queue<int> q;
        for (int j: vec[color[i]]) {
            q.push(j);
            vis[j] = true;
        }
        int cnt = 0;
        while (!q.empty()) {
            tot++;
            assert(tot <= n);
            int u = q.front(); q.pop();
            if (i == u) continue;
            if (vis[jump[0][u]]) continue;
            cnt++;
            for (int j: vec[color[jump[0][u]]]) {
                q.push(j);
                vis[j] = true;
            }
        }
        mins(ans, cnt);
    }
    cout << ans;
}

int32_t main() {
    clock_t startTime = clock();
    cin.tie(0)->sync_with_stdio(false);
    srand(time(0));

    std::istream &in(std::cin);
    std::ostream &out(std::cout);

    int32_t queries = 1;

#ifdef test_cases
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
    cin >> queries;
#else
    // cin >> queries;
#endif

    for (int32_t test_case = 1; test_case <= queries; test_case++) {
        print(test_case);
        solution(cin, cout, test_case);
        cout << nl;
    }

#ifdef sunnatov
    cout << "Time: " << (int) ((double) (clock() - startTime) / CLOCKS_PER_SEC * 1000) << " ms" << endl;
#endif

    return EXIT_SUCCESS;
}

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