Submission #838186

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

uzastopni

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

#define int long long

template<typename T>
string tostr(const T& value) {
    ostringstream oss;
    oss << value;
    return oss.str();
}

template<typename... Args>
string fstr(const string& format, Args... args) {
    string result = format;
    size_t pos = 0;
    size_t argIndex = 0;

    auto replaceArg = [&](const auto& arg) {
        pos = result.find("{}", pos);
        if (pos != string::npos) {
            result.replace(pos, 2, tostr(arg));
            ++argIndex;
        }
    };

    (replaceArg(args), ...);

    return result;
}

/*
 * Tree
 */
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);
        }
    }
};

/*
 * LCA
 * LCA-related functions

*/
template<int SZ, int LOG> struct LCA {
    int n;
    int bl[SZ][LOG];
    int dpth[SZ];
    void init_dfs(int nd, int par, vector<int>* adj) {
        dpth[nd] = dpth[par]+1;
        bl[nd][0] = par;
        for (int i = 1; i < LOG; i++)
            bl[nd][i] = bl[bl[nd][i-1]][i-1];
        for (auto i : adj[nd])
            if (i != par)
                init_dfs(i, nd, adj);
    }
    void init(int N, Tree<SZ>& tree, int R = 0) {
        n = N;
        init_dfs(R, R, tree.adj);
    }
    int jmp(int nd, int k) {
        for (int i = LOG-1; i >= 0; i--) if ((k >> i) & 1) nd = bl[nd][i];
        return nd;
    }
    int lca(int a, int b) {
        if (dpth[a] < dpth[b]) swap(a, b);
        a = jmp(a, dpth[a] - dpth[b]);
        if (a == b) return a;
        for (int i = LOG-1; i >= 0; i--) {
            int A = bl[a][i], B = bl[b][i];
            if (A != B) a = A, b = B;
        }
        return bl[a][0]; 
    }
    bool onpath(int a, int b, int c) {
        // whether c is on path from a --> b
        int lc = lca(a, b);
        return (dpth[lc] <= dpth[c]) && ((dpth[c] <= dpth[a] && jmp(a, dpth[a] - dpth[c]) == c) || (dpth[c] <= dpth[b] && jmp(b, dpth[b] - dpth[c]) == c));
    }
};

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]];

    // cout << "----------------- " << nd << endl;
    // for (int i = 0; i < 4; i++) {
    //     cout << fstr("i: {}, l ", i);
    //     for (int j = 0; j < 4; j++)
    //         cout << l[nd][i][j];
    //     cout << endl;
    //     cout << fstr("i: {}, r ", i);
    //     for (int j = 0; j < 4; j++)
    //         cout << r[nd][i][j];
    //     cout << endl;
    // }

    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