Submission #1058772

#	Time	Username	Problem	Language	Result	Execution time	Memory
1058772		gyg	Cat in a tree (BOI17_catinatree)	C++17	11 / 100	19 ms	20828 KiB

catinatree

#include <bits/stdc++.h>
using namespace std;
#define int long long
using pii = pair<int, int>;
const int MAX_N = 1500 + 5;

int n, k;
array<vector<int>, MAX_N> adj;

/************************************/
array<array<int, MAX_N>, MAX_N> floyd;
array<bool, MAX_N> seen;
void dfs(int u, int orig, int d = 0) {
    seen[u] = true, floyd[orig][u] = d;
    for (int v : adj[u])
        if (!seen[v]) dfs(v, orig, d + 1);
}
void precomp() {
    for (int u = 1; u <= n; u++) {
        fill(seen.begin(), seen.end(), false);
        dfs(u, u);
    }
}
/************************************/


array<int, MAX_N> siz;
array<bool, MAX_N> del;
void update_siz(int u, int par = -1) {
    siz[u] = 1;
    for (int v : adj[u]) {
        if (v == par) continue;
        if (del[v]) continue;
        update_siz(v, u), siz[u] += siz[v];
    }
}
int find_cent(int u, int tot_siz, int par = -1) {
    for (int v : adj[u]) {
        if (v == par) continue;
        if (del[v]) continue;
        if (siz[v] * 2 > tot_siz) return find_cent(v, tot_siz, u);
    }
    return u;
}

array<vector<int>, MAX_N> ancs; // Reversed
array<unordered_map<int, int>, MAX_N> dist, subcent;
array<unordered_map<int, vector<pii>>, MAX_N> dists; // [itself] = whole thing, everything gets reversed after queries
void find_dists(int u, int d, int par, int cent, int sub) {
    dist[cent][u] = d, dists[cent][sub].push_back({d, u});
    subcent[cent][u] = sub;
    ancs[u].push_back(cent);
    for (int v : adj[u]) {
        if (v == par) continue;
        if (del[v]) continue;
        find_dists(v, d + 1, u, cent, sub);
    }
}

void build(int u = 1) {
    update_siz(u);
    int cent = find_cent(u, siz[u]);
    
    dist[cent][cent] = 0, dists[cent][cent].push_back({0, cent});
    ancs[cent].push_back(cent);
    for (int v : adj[cent]) {
        if (del[v]) continue;
        find_dists(v, 1, cent, cent, v);
        sort(dists[cent][v].begin(), dists[cent][v].end());
        dists[cent][cent].insert(dists[cent][cent].begin(), dists[cent][v].begin(), dists[cent][v].end());
    }
    sort(dists[cent][cent].begin(), dists[cent][cent].end());

    del[cent] = true;
    for (int v : adj[cent])
        if (!del[v]) build(v);
}

int mini_query(int u, int d, int exc) {
    int i = lower_bound(dists[u][u].begin(), dists[u][u].end(), make_pair(d, -1ll)) - 1 - dists[u][u].begin(); // First index with dist < d
    int j = (exc == -1) ? -1 : lower_bound(dists[u][exc].begin(), dists[u][exc].end(), make_pair(d, -1ll)) - 1 - dists[u][exc].begin();
    int ans = i + 1 - (j + 1);
    return ans;
}
int query(int u) {
    reverse(ancs[u].begin(), ancs[u].end());
    int ans = 0;
    for (int i = 0; i < ancs[u].size(); i++) {
        int v = ancs[u][i];
        int new_ans = (i == 0) ? mini_query(v, k, -1) : mini_query(v, k - dist[v][u], subcent[v][u]);
        ans += new_ans;
    }
    // cout << u << ": ";
    // for (int v : ancs[u]) cout << v << " ";
    // cout << endl;

    // cout << u << ": " << ans << endl;

    return ans;
}

vector<pii> order;
array<bool, MAX_N> used;
void update_used(int u) {
    // for (int v : ancs[u]) {
    //     int d = k - dist[v][u];
    //     while (dists[v][v].size() && (dists[v][v].back().first < d || used[dists[v][v].back().second])) { // sus
    //         used[dists[v][v].back().second] = true;
    //         dists[v][v].pop_back();
    //     } 
    // }
    for (int v = 1; v <= n; v++)
        if (floyd[u][v] < k) used[v] = true;
}

int comp() {
    for (int u = 1; u <= n; u++) order.push_back({query(u), u}); 
    sort(order.begin(), order.end());   
    for (int u = 1; u <= n; u++) reverse(dists[u][u].begin(), dists[u][u].end());  
    
    int ans = 0;
    for (pii x : order) {
        int u = x.second;
        if (used[u]) continue;
        ans++;
        update_used(u);

        // cout << "DELETED: " << u << endl;
        // cout << "LEFT: ";
        // for (int v = 1; v <= n; v++)
        //     if (!used[v]) cout << v << " ";
        // cout << endl;
    }
    return ans;
}

signed main() {
    // freopen("cat.in", "r", stdin);
    cin >> n >> k; assert(n <= 1500);
    for (int u = 2; u <= n; u++) {
        int v; cin >> v; v++;
        adj[u].push_back(v), adj[v].push_back(u);
    }

    precomp();
    build();

    // for (int u = 1; u <= n; u++) {
    //     cout << u << ": ";
    //     for (pii x : dists[u][u]) cout << x.second << "," << x.first << " ";
    //     cout << endl; 
    // }
 
    int ans = comp();
    cout << ans << '\n';
}

Compilation message (stderr)

catinatree.cpp: In function 'long long int query(long long int)':
catinatree.cpp:88:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   88 |     for (int i = 0; i < ancs[u].size(); i++) {
      |                     ~~^~~~~~~~~~~~~~~~

• Subtask 1: 11 / 11
• Subtask 2: 0 / 40
• Subtask 3: 0 / 49