Submission #838189

#	Time	Username	Problem	Language	Result	Execution time	Memory
838189		Shreyan_Paliwal	Uzastopni (COCI15_uzastopni)	C++17	160 / 160	7 ms	7892 KiB

uzastopni

// #pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;

#define int long long

template<int SZ> struct Tree {
    int n;
    vector<int> adj[SZ];
    void init(int N) {
        n = N;
        for (int i = 0; i < n-1; i++) {
            int a, b; cin >> a >> b; a--, b--;
            adj[a].push_back(b);
            adj[b].push_back(a);
        }
    }
};

const int maxn = 10001;
const int maxk = 101;

int n;
int v[maxn];
int vis[maxk];
Tree<maxn> tree;
bitset<maxk> l[maxn][maxk], r[maxn][maxk];

void dfs(int nd, int par) {
    if (vis[v[nd]]) return;
    vis[v[nd]] = 1;

    for (auto i : tree.adj[nd])
        dfs(i, nd);
        

    // set l
    l[nd][v[nd]][v[nd]] = 1;
    for (int i = v[nd]+1; i < maxk; i++)
        if (l[nd][v[nd]][i-1])
            for (auto j : tree.adj[nd])
                if (j != par)
                    l[nd][v[nd]] |= l[j][i];         

    // set r
    r[nd][v[nd]][v[nd]] = 1;
    for (int i = v[nd]-1; i >= 0; i--)
        if (r[nd][v[nd]][i+1])
            for (auto j : tree.adj[nd])
                if (j != par)
                    r[nd][v[nd]] |= r[j][i];

    // set rest
    for (int i = 0; i < maxk; i++)
        if (i != v[nd] && l[nd][v[nd]][i])
            r[nd][i] = r[nd][v[nd]];

    for (int i = 0; i < maxk; i++)
        if (i != v[nd] && r[nd][v[nd]][i])
            l[nd][i] = l[nd][v[nd]];

    vis[v[nd]] = 0;
}


// #define LOCAL
// #define CODEFORCES
signed main() {
    #ifndef LOCAL
    cin.tie(nullptr) -> ios::sync_with_stdio(false);
    #endif
    #ifdef LOCAL
    freopen("main.in", "r", stdin);
    #endif
    
    cin >> n;
    for (int i = 0; i < n; i++) {cin >> v[i]; v[i]--;}
    tree.init(n);

    dfs(0, 0);

    int ans = 0;
    for (int i = 0; i < maxk; i++)
        ans += l[0][i].count();
    cout << ans << endl;

    return 0;
}

• Subtask 1: 160 / 160