読者です 読者をやめる 読者になる 読者になる

SRM435 Div1Easy(Div2Med) CellRemoval

問題概要
vector parent が与えられる。parent[i] とは i 番目のノードの親が parent[i] であることを示す。deletedCell で削除するノードが与えられたとき、残ったノードのうち葉のノードの数をカウントせよ。

解法
子を参照する形の木を、与えられた vector parent から生成し、削除するノードで打ち切って dfs していけばよい。
vector > tree; に対して、tree[ parent[i] ].push_back(i); を for文で回すと木が生成できる。ただし、木全体の親ノードに対しては parent[i] = -1 となるので、treeにノードを追加する操作はない。代わりにそのときの i を木の走査のスタート位置として保持しておく。

#define PB push_back
#define REP(i,a,b) for(int i=a;i<b;i++)
#define rep(i,n) REP(i,0,n)

typedef vector<int> VI;
typedef vector<VI> VVI;

class CellRemoval {
public:
  
  VVI tree;
  int del;
  
  int dfs(int idx) {
    
    if(idx == del) {
      return 0;
    }
    
    if(tree[idx].empty()) {
      return 1;
    }
    
    int res = 0;
    rep(i, tree[idx].size()) {
      res += dfs(tree[idx][i]);
    }
    
    return res;
  }  
  
  int cellsLeft(vector <int> parent, int deletedCell) {
    
    int const N = parent.size();
    del = deletedCell;
    tree.clear();
    tree.resize(N);

    int start;
    rep(i, N) {
      if(parent[i] == -1) start = i;
      else tree[parent[i]].PB(i);
    }
    
    return dfs(start);
  }
};