/ SeriousOJ /

Record Detail

Wrong Answer


  
# Status Time Cost Memory Cost
#1 Accepted 41ms 38.781 MiB
#2 Accepted 40ms 38.785 MiB
#3 Wrong Answer 63ms 38.785 MiB
#4 Wrong Answer 63ms 38.863 MiB

Code

#include <bits/stdc++.h>
#define endl '\n'
using namespace std;
#define ll long long
#define mod 1000000007
#define int long long
#define all(x)      x.begin(),x.end()
#define allr(x)     x.rbegin(),x.rend()
#define CheckBit(x,k)   (x & (1LL << k))
#define SetBit(x,k)     (x |= (1LL << k))
#define ClearBit(x,k)   (x &= ~(1LL << k))
#define LSB(mask)       __builtin_ctzll(mask)
#define MSB(mask)       63-__builtin_clzll(mask) 
#define print(x)    cout << #x << " : " << x << endl
#define error1(x)   cerr << #x << " = " << (x) <<endl
#define coutall(v)  for(auto &it: v) cout<<it<<' '; cout<<endl
#define Abid_52     ios::sync_with_stdio(false);cin.tie(0),cin.tie(0)
#define error2(a,b) cerr<<"( "<<#a<<" , "<<#b<<" ) = ( "<<(a)<<" , "<<(b)<<" )\n"
#define UNIQUE(x)   sort(all(x)), x.erase(unique(all(x)), x.end()), x.shrink_to_fit()
template <typename T, typename U> T ceil(T x, U y) {return (x > 0 ? (x + y - 1) / y : x / y);}
template <typename T, typename U> T floor(T x, U y) {return (x > 0 ? x / y : (x - y + 1) / y);}
const int N = 2e5 + 10;
const int Log = 20;
vector<int> g[N + 1];
int timer = 0;
int tin[N + 1];  // for node entering time
int tout[N + 1]; // for node existing time
int table[N + 1][Log + 1];
int level[N + 1];
int apple[N + 5];
vector<int>vis(N + 5, 0);
int track = -1;
int ans = 0;
void reset()
{
    for (int i = 0; i < N; i++)
    {
        g[i].clear();
        tin[i] = 0;
        tout[i] = 0;
        level[i] = 0;
        vis[i] = 0;
        for (int j = 0; j < Log; j++)
        {
            table[i][j] = -1;
        }
    }
    timer = 0;
    track = -1;
    ans = 0;

}
void dfs(int node, int par, int lev = 0)
{
    tin[node] = ++timer; // node entering time
    table[node][0] = par;
    level[node] = lev;
    for (int i = 1; i < Log; i++)
    {
        int p = table[node][i - 1];
        table[node][i] = table[p][i - 1];
    }
    for (int child : g[node])
    {
        if (child != par)
        {
            dfs(child, node, lev + 1);
        }
    }
    tout[node] = ++timer; // node existing time
}
bool isAncestor(int u, int v) // check if u is ancestor of v or not
{
    return (tin[u] <= tin[v] && tout[u] >= tout[v]);
}
int lca(int u, int v)
{
    if (isAncestor(u, v))
    {
        return u;
    }
    if (isAncestor(v, u))
    {
        return v;
    }
    for (int i = Log - 1; i >= 0; i--)
    {
        if (!isAncestor(table[u][i], v))
        {
            u = table[u][i];
        }
    }
    return table[u][0];
}
int distance(int u, int v) // distance between two node u and v
{
    int l = lca(u, v);
    return level[u] + level[v] - (level[l] << 1); // level[u]+level[v]-2*level[lca(u,v)]
}
int kth(int u, int k)
{
    for (int i = 0; i <= Log; i++)
        if (k & (1 << i))
            u = table[u][i];
    return u;
}
int findKth(int u, int v, int k) // kth node from u to v, 0th node is u
{
    int l = lca(u, v);
    int d = level[u] + level[v] - (level[l] << 1);
    if (level[l] + k <= level[u])
        return kth(u, k);
    k -= level[u] - level[l];
    return kth(v, level[v] - level[l]-k);
}
void dfs1(int node){
    // print(node);
    vis[node] = 1;
    if(track == -1)track = node;
    if(track != -1 && apple[node] == 1){
        // print(node);
        // print(track);
        ans += (distance(node, track));
        track = node;
    }
    for(auto it : g[node]){
        if(!vis[it]){
            dfs1(it);
        }
    }
}
void solve()
{
    int n;
    cin >> n;
    int app = -1;
    for(int i = 1; i <= n; i++){
        cin >> apple[i];
        if(apple[i] == 1 && app == -1){
            app = i;
        }
    }
    for (int i = 1; i < n; i++)
    {
        int u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    if(app == -1){
        cout << 0 << endl;
        return;
    }
    dfs(1, 1);
    dfs1(app);
    // for (int i = 1; i <= n; i++)
    // {
    //     for (int j = 0; j <= 10; j++)
    //     {
    //         cout << table[i][j] << " ";
    //     }
    //     cout << endl;
    // }
    // int q;
    // cin >> q;
    // while (q--)
    // {
    //     int a, b;
    //     cin >> a >> b;
    //     cout << lca(a, b) << endl;
    //     // cout << distance(a, b) << endl;
    // }
    // print(distance(1, 8));
    cout << ans << endl;
}

int32_t main()
{
    Abid_52;
    int t = 1;
    cin >> t;
    for (int tc = 1; tc <= t; ++tc)
    {
        // cout << "Case " << tc << ":";
        reset();
        solve();
    }
}

Information

Submit By
Type
Submission
Problem
P1078 Apple on Tree
Contest
Bangladesh 2.0
Language
C++20 (G++ 13.2.0)
Submit At
2024-08-16 16:48:33
Judged At
2024-11-11 03:13:52
Judged By
Score
4
Total Time
63ms
Peak Memory
38.863 MiB