Submission #738358

#	Time	Username	Problem	Language	Result	Execution time	Memory
738358		JakobZorz	Cat in a tree (BOI17_catinatree)	C++14	100 / 100	210 ms	34424 KiB

catinatree

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <set>
#include <stack>
#include <limits.h>
 
#define ll long long
using namespace std;
const int MOD = 1e9 + 7;
 
//#define SEGTREE
#define TREE
//#define DSU
 
#ifdef SEGTREE
template<class Type>
class SegmentTree {
    Type (*func)(Type a, Type b);
    vector<Type> nodes;
    vector<int> l;
    vector<int> r;
    int size_log2;
    Type neutral;
    
    void init_node(int node) {
        if(node >= (1 << size_log2))
            return;
        
        l[2 * node] = l[node];
        r[2 * node] = (l[node] + r[node]) / 2;
        init_node(2 * node);
        
        l[2 * node + 1] = (l[node] + r[node]) / 2;
        r[2 * node + 1] = r[node];
        init_node(2 * node + 1);
        
        nodes[node] = func(nodes[2 * node], nodes[2 * node + 1]);
    }
    
    void update_node(int node) {
        if(node < (1 << size_log2))
            nodes[node] = func(nodes[2 * node], nodes[2 * node + 1]);
        
        if(node != 1)
            update_node(node / 2);
    }
    
    Type get(int node, int ll, int rr) {
        if(rr <= l[node] || ll >= r[node])
            return neutral;
        
        if(ll <= l[node] && rr >= r[node])
            return nodes[node];
        
        Type left = get(2 * node, ll, rr);
        Type right = get(2 * node + 1, ll, rr);
        
        return func(left, right);
    }
    
public:
    SegmentTree(Type (*func)(Type a, Type b), vector<Type> vector, Type neutral) : func(func), neutral(neutral) {
        size_log2 = 0;
        while((1 << size_log2) < vector.size())
            size_log2++;
        
        nodes.resize((1 << size_log2) * 2);
        l.resize((1 << size_log2) * 2);
        r.resize((1 << size_log2) * 2);
        
        for(int i = 0; i < vector.size(); i++)
            nodes[(1 << size_log2) + i] = vector[i];
        
        l[1] = 0;
        r[1] = 1 << size_log2;
        init_node(1);
    }
    
    void set(int idx, Type el) {
        nodes[(1 << size_log2) + idx] = el;
        update_node((1 << size_log2) + idx);
    }
    
    Type get(int l, int r) {
        return get(1, l, r);
    }
    
    Type get(int idx) {
        return nodes[(1 << size_log2) + idx];
    }
    
    Type get() {
        return nodes[1];
    }
    
    int get_first_larger_or_eq_than(int bound) {
        return get_first_larger_or_eq_than(1, bound);
    }
    
    int get_first_larger_or_eq_than(int node, int bound) {
        if(node >= (1 << size_log2)) {
            return node - (1 << size_log2);
        }
        
        if(nodes[node * 2] > bound)
            return get_first_larger_or_eq_than(node * 2, bound);
        else
            return get_first_larger_or_eq_than(node * 2 + 1, bound - nodes[node * 2]);
    }
};
 
#endif
 
#ifdef TREE
#define POW 18
 
struct Node {
    int parents[POW];
    //int max_weights[POW];
    vector<int> conns;
    //vector<int> weights;
    int depth;
    //int subtree_depth;
};
 
Node nodes[200000];
int n;
 
void setup(int node, int depth) {
    nodes[node].depth = depth;
    //for(int i = 1; i < POW; i++) {
        //nodes[node].parents[i] = nodes[nodes[node].parents[i - 1]].parents[i - 1];
        //nodes[node].max_weights[i] = max(nodes[nodes[node].parents[i - 1]].max_weights[i - 1], nodes[node].max_weights[i - 1]);
    //}
    
    for(int i = 0; i < nodes[node].conns.size(); i++) {
        int child = nodes[node].conns[i];
        //int weight = nodes[node].weights[i];
        if(child != nodes[node].parents[0]) {
            nodes[child].parents[0] = node;
            //nodes[child].max_weights[0] = weight;
            setup(child, depth + 1);
        }
    }
}
 
void setup_tree(int root) {
    nodes[root].parents[0] = root;
    setup(root, 0);
}
 
/*int get_subtree_depth(int node) {
    if(nodes[node].subtree_depth)
        return nodes[node].subtree_depth;
    
    int max_depth = 1;
    for(int child : nodes[node].conns) {
        if(child == nodes[node].parents[0])
            continue;
        max_depth = max(max_depth, get_subtree_depth(child) + 1);
    }
    nodes[node].subtree_depth = max_depth;
    return max_depth;
}*/
 
/*int get_parent(int node, int depth) {
    if(depth > nodes[node].depth)
        return -1;
    
    int climber = node;
    for(int i = 0; i < POW; i++) {
        if(depth & (1 << i) && climber != -1)
            climber = nodes[climber].parents[i];
    }
    
    return climber;
}
 
int lca(int a, int b) {
    if(nodes[a].depth < nodes[b].depth)
        swap(a, b);
    
    a = get_parent(a, nodes[a].depth - nodes[b].depth);
    
    if(a == b)
        return a;
    
    for(int i = POW - 1; i >= 0; i--) {
        if(nodes[a].parents[i] != nodes[b].parents[i]) {
            a = nodes[a].parents[i];
            b = nodes[b].parents[i];
        }
    }
    
    return nodes[a].parents[0];
}*/
 
#endif
 
#ifdef DSU
 
class Dsu {
    vector<int> arr;
    int num_sets;
    
public:
    Dsu(int size) {
        arr = vector<int>(size, -1);
        num_sets = size;
    }
    
    int merge(int a, int b) {
        a = get(a);
        b = get(b);
        
        if(a == b)
            return a;
        
        if(arr[a] > arr[b])
            swap(a, b);
        
        arr[a] += arr[b];
        arr[b] = a;
        num_sets--;
        return a;
    }
    
    int get(int a) {
        if(arr[a] < 0)
            return a;
        arr[a] = get(arr[a]);
        return get(arr[a]);
    }
    
    int get_size(int a) {
        return -arr[get(a)];
    }
    
    int get_num_sets() {
        return num_sets;
    }
};
 
#endif

int k;
int c[200000];
int nodes_down[200000];

bool cmp(int a, int b) {
    return nodes[a].depth > nodes[b].depth;
}

void color(int node, int col) {
    if(c[node] >= col)
        return;
    
    c[node] = col;
    for(int conn : nodes[node].conns)
        color(conn, col - 1);
}

int main() {
    ios::sync_with_stdio(false);
    cout.tie(NULL);
    cin.tie(NULL);
    
    cin >> n >> k;
    
    for(int i = 0; i < n - 1; i++) {
        int a, b = i + 1;
        cin >> a;// >> b;
        //a--;
        //b--;
        nodes[a].conns.push_back(b);
        nodes[b].conns.push_back(a);
    }
    
    setup_tree(0);
    
    for(int i = 0; i < n; i++)
        nodes_down[i] = i;
    
    sort(nodes_down, nodes_down + n, cmp);
    
    for(int i = 0; i < n; i++) {
        if(c[nodes_down[i]] == 0) {
            color(nodes_down[i], k);
        }
    }
    
    int cnt = 0;
    for(int i = 0; i < n; i++)
        if(c[i] == k)
            cnt++;
    
    cout << cnt << endl;
    /*for(int i = 0; i < n; i++)
        if(c[i] == k)
            cout << i + 1 << " ";
    cout << endl;*/
    
    return 0;
}

Compilation message (stderr)

catinatree.cpp: In function 'void setup(int, int)':
catinatree.cpp:138:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  138 |     for(int i = 0; i < nodes[node].conns.size(); i++) {
      |                    ~~^~~~~~~~~~~~~~~~~~~~~~~~~~

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