(15899) 트리와 색깔

https://www.acmicpc.net/problem/15899

단계별로 풀어보기

52단계(세그먼트 트리 2) 10번째

 

 

 

오일러 경로 테크닉과 머지 소트 트리를 결합하는 문제. 오일러 경로 테크닉에서 세그먼트 트리의 초기 단일 값은 트리에서 초기 정점 인덱스가 아닌 전위 순회로 방문한 순서라는 새로운 인덱스로 설정해야 한다는 것만 유의하면 된다.

 

 

 

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
using namespace std;
#define ALL(v) (v).begin(), (v).end()
#define p 1000000007

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  
  int n, m, c, cnt= 0, ans= 0;
  
  cin >> n >> m >> c;
  
  vector<int> v0(n+1), v(n+1), in(n+1), out(n+1);
  vector<vector<int>> adj(n+1);
  
  for(int i=1;i<=n;i++)  cin >> v0[i];
  
  for(int i=1;i<n;i++){
    int a, b;
    
    cin >> a >> b;
    
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  
  auto dfs= [&](auto self, int now, int prev) -> void {
    in[now]= ++cnt;
    v[cnt]= v0[now];
    
    for(int next:adj[now]){
      if(next==prev)  continue;
      
      self(self, next, now);
    }
    
    out[now]= cnt;
  };
  
  cnt= 0;
  dfs(dfs, 1, 0);
  
  int treesize= (1<<(((int)ceil(log2(n+1)))+1));
  vector mstree(treesize, vector<int> ());
  
  auto init= [&](auto self, int s, int e, int node) -> void {
    if(s==e){
      mstree[node].push_back(v[s]);
    }else{
      self(self, s, s+(e-s)/2, node*2);
      self(self, s+(e-s)/2+1, e, node*2+1);
      
      mstree[node].resize(mstree[node*2].size()+mstree[node*2+1].size());
      merge(ALL(mstree[node*2]), ALL(mstree[node*2+1]), mstree[node].begin());
    }
  };
  
  auto query= [&](auto self, int s, int e, int node, int l, int r, int k) -> int {
    if(r<s||e<l)  return 0;
    
    if(l<=s&&e<=r){
      return  upper_bound(ALL(mstree[node]), k)-mstree[node].begin();
    }else{
      long lsum= self(self, s, s+(e-s)/2, node*2, l, r, k);
      long rsum= self(self, s+(e-s)/2+1, e, node*2+1, l, r, k);
      return lsum+rsum;
    }
  };
  
  init(init, 1, n, 1);
  
  while(m--){
    int vi, ci;
    
    cin >> vi >> ci;
    
    ans= (ans+query(query, 1, n, 1, in[vi], out[vi], ci))%p;
  }
  
  cout << ans;
  
  return 0;
}