Submission #12141884


Source Code Expand

Copy
#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#define pi 3.141592653589793238
#define int long long
using namespace __gnu_pbds;
using namespace std;
template <typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

int st[200005];
int en[200005];
int timer=0;
vector<vector<int>>adj;
int sz[200005];
int col[200005];
int ans[200005];
int node[800005];
set<pair<int,pair<int,int> > >s;

int dfs(int ver,int pr)
{
  st[ver]=++timer;
  node[timer]=ver;
  for(auto i:adj[ver])
  {
    if(i!=pr)
      sz[ver]+=dfs(i,ver);
  }
  en[ver]=++timer;
  sz[ver]++;
  return sz[ver];
}
int n;
void madar()
{
  for(int cc=1;cc<=n;cc++)
  {
    int tot=0;
    auto j=s.lower_bound({cc,{0,-1}});
    vector<pair<int,pair<int,int> > >temp;
    while (j!=s.end())
    {
      if(j->first!=cc)
        break;
      if(j->second.second>1e9)
        break;
      temp.push_back(*j);
      j++;
    }    
    tot+=sz[1];
    for(auto j:temp)
    {
      tot-=sz[node[j.second.first]];
      s.erase(j);
    }     
    ans[cc]+=(tot*(tot+1))/2;  
  }
}

void dfs2(int ver,int pr)
{
  for(auto i:adj[ver])
  {
    if(i!=pr)
    {
      dfs2(i,ver);
      int tot=0;
      auto j=s.lower_bound({col[ver],{st[ver],-1}});
      vector<pair<int,pair<int,int> > >temp;
      while (j!=s.end())
      {
        if(j->first!=col[ver])
          break;
        if(j->second.second>en[ver])
          break;
        temp.push_back(*j);
        j++;
      }     
      tot+=sz[i];
      for(auto j:temp)
      {
        tot-=sz[node[j.second.first]];
        s.erase(j);
      }     
      ans[col[ver]]+=(tot*(tot+1))/2;  
    }
  }
  s.insert({col[ver],{st[ver],en[ver]}});
}

signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);  
    cout.tie(0);   
#ifndef ONLINE_JUDGE
    if(fopen("INPUT.txt","r"))
    {
    freopen ("INPUT.txt" , "r" , stdin);
    freopen ("OUTPUT.txt" , "w" , stdout);
    }
#endif

// -------------------------------------Code starts here---------------------------------------------------------------------     

  cin>>n;
  adj.resize(n+1);
  for(int i=1;i<=n;i++)
    cin>>col[i];
  int x,y;
  for(int i=0;i<n-1;i++)
  {
    cin>>x>>y;
    adj[x].push_back(y);
    adj[y].push_back(x);
  }
  dfs(1,0);
  dfs2(1,0);
  madar();
  int fi=(n*(n+1))/2;
  for(int i=1;i<=n;i++)
  {
    cout<<(fi-ans[i])<<'\n';
  }
 
// -------------------------------------Code ends here------------------------------------------------------------------
    clock_t clk;
  clk = clock();

    clk = clock() - clk;
  cerr << fixed << setprecision(6) << "Time: " << ((double)clk)/CLOCKS_PER_SEC << "\n";
  return 0;
}

Submission Info

Submission Time
Task F - path pass i
User babayaga
Language C++ (GCC 9.2.1)
Score 600
Code Size 2854 Byte
Status AC
Exec Time 351 ms
Memory 68816 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 600 / 600
Status
AC × 5
AC × 24
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All line_01.txt, line_3_01.txt, random_01.txt, random_02.txt, random_03.txt, random_100_01.txt, random_100_02.txt, random_10_01.txt, random_10_02.txt, random_1_01.txt, random_1_02.txt, random_2_01.txt, random_2_02.txt, random_2_03.txt, random_3_01.txt, random_3_02.txt, random_3_03.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt, star_01.txt, star_3_01.txt
Case Name Status Exec Time Memory
line_01.txt AC 236 ms 68816 KB
line_3_01.txt AC 114 ms 60884 KB
random_01.txt AC 340 ms 38720 KB
random_02.txt AC 336 ms 38764 KB
random_03.txt AC 351 ms 38764 KB
random_100_01.txt AC 278 ms 29128 KB
random_100_02.txt AC 279 ms 29524 KB
random_10_01.txt AC 241 ms 26364 KB
random_10_02.txt AC 226 ms 26368 KB
random_1_01.txt AC 190 ms 26364 KB
random_1_02.txt AC 194 ms 26276 KB
random_2_01.txt AC 220 ms 26232 KB
random_2_02.txt AC 219 ms 26272 KB
random_2_03.txt AC 218 ms 26176 KB
random_3_01.txt AC 213 ms 26408 KB
random_3_02.txt AC 224 ms 26300 KB
random_3_03.txt AC 227 ms 26816 KB
sample_01.txt AC 5 ms 3756 KB
sample_02.txt AC 2 ms 3764 KB
sample_03.txt AC 5 ms 3760 KB
sample_04.txt AC 2 ms 3800 KB
sample_05.txt AC 3 ms 3712 KB
star_01.txt AC 220 ms 39296 KB
star_3_01.txt AC 136 ms 37960 KB