Submission #128419

#	Time	Username	Problem	Language	Result	Execution time	Memory
128419		Gustav	Uzastopni (COCI15_uzastopni)	C++14	160 / 160	36 ms	6136 KiB

uzastopni

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pi;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<pi> vpi;
typedef vector<vi> vvi;
const int inf = 0x3f3f3f3f;
const ll linf = 123456789012345678;
const ll mod = 1000000007;
#pragma GCC optimize("Ofast")
#define all(x) x.begin(), x.end()
#define debug(x) cerr << #x << " = " << x << endl
#define sz(x) ((int)(x).size())

int n;
vvi adj;
vi k;

vvi adj2;
vi visited;

void dfs2(int node, bool back){
    if(visited[node]==1) return;
    visited[node] = 1;
    if((!back && node < 60000) || (back && (node >= 60000 || node < 40000))){
    for(int x : adj2[node]){
        dfs2(x, back);
    }
    }
    if(!back && node >= 60000) dfs2(node-20000+1, back);
    if(back && node >= 40000 && node < 60000) dfs2(node+20000-1, back);
}

pair<vi, vi> dfs(int node, int par){
    vector<pair<vi, vi>> results;
    for(int x : adj[node]){
        if(x == par) continue;
        results.push_back(dfs(x, node));
    }

    for(int i = 0; i < sz(results); i++){
        for(int x : results[i].first){
            adj2[x+40000].push_back(i);
            adj2[i+20000].push_back(x+40000);
        }
        for(int x : results[i].second){
            adj2[x+60000].push_back(i+20000);
            adj2[i].push_back(x+60000);
        }
    }
    dfs2(60000 + k[node], false);
    dfs2(40000 + k[node], true);
    vi start, end;
    for(int i = 0; i <= k[node]; i++){
        if(visited[40000+i]) start.push_back(i);
    }
    for(int i = k[node]; i < 100; i++){
        if(visited[60000+i]) end.push_back(i);
    }
    for(int i = 0; i < sz(results); i++){
        for(int x : results[i].first){
            adj2[x+40000].clear();
            adj2[i+20000].clear();
            visited[i] = 0;
            visited[x+40000] = 0;
            visited[x+60000-1] = 0;
        }
        for(int x : results[i].second){
            adj2[x+60000].clear();
            adj2[i].clear();
            visited[i+20000] = 0;
            visited[x+60000] = 0;
            visited[x+40000+1] = 0;
        }
    }
    visited[k[node]+60000] = 0;
    visited[k[node]+40000+1] = 0;
    visited[k[node]+40000] = 0;
    visited[k[node]+60000-1] = 0;
    return {start, end};
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    adj2 = vvi(100000);
    visited = vi(100000);

    cin >> n;
    k = vi(n);
    adj = vvi(n);
    for(int i = 0; i < n; i++) {
        cin >> k[i];
        k[i]--;
    }
    for(int i = 0; i < n-1; i++){
        int u, v;
        cin >> u >> v;
        u--, v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    pair<vi, vi> p = dfs(0, -1);
    cout << sz(p.first)*sz(p.second) << endl;
}

• Subtask 1: 160 / 160