ABC198 E 解答

E - Unique Color(1161)

問題文

N 頂点からなる木が与えられます。
i 番目の辺は頂点 Ai と頂点 Bi をつないでいます。頂点 i は色 Ci で塗られています (この問題では、色は整数として表されます)。
頂点 1 から頂点 x への最短パスに、頂点 x と同じ色で塗られた頂点が頂点 x 以外に存在しないとき、頂点 x は よい頂点 であるといいます。
よい頂点を全て求めてください。

制約

2 ≤ N ≤ 10^5
1 ≤ Ci ≤ 10^5
1 ≤ Ai, Bi ≤ N
与えられるグラフは木である
入力は全て整数

考察

まず、計算量を無視して考えていきます。

良い頂点は

頂点１からある頂点 x までの経路で頂点 x と同じ色がない

というのが条件です。ですので、DFSやBFSなどで頂点1 と x までの経路を見てあげて、その経路にxと同じ色がなければ良い頂点となり、それを全ての頂点に対して行えば答えが求まります。

ただこれだと間に合いませんので工夫をしましょう。

今回は、現時点で通った頂点の色を記録していきながらDFSで探索をしていきます。これにより、自分よりも上の頂点に書かれている色を管理することができますので、１回の探索で全ての頂点の判定を行うことが可能となります。

本手法ではグラフとなる木に加えて、どこからでも書き込める掲示板のような配列を用意して探索をしていきます。この掲示板には今まで通った頂点に書かれてた色とその数を記録していきます。

実際の探索のやり方を説明していきます。

まず、その頂点に到着しましたら掲示板を見ます。もし掲示板に頂点の色が書き込まれていなければ、その頂点は良い頂点と見なせます。一方で、既に書かれていましたら残念ながら良い頂点ではありません。

判定が終わったら、次の頂点に進む前に掲示板を更新しましょう。その頂点に書かれている色の数を+1してあげます。

こんな感じで、どんどん下っていきます。

さて、探索を続けて行ったところあなたは一番下層にたどり着きました。ということで、今度は来た道を戻りましょう。そのさいに、戻り道の頂点に書かれている色に関して掲示板を更新しましょう。掲示板にはその頂点よりも上にある色の情報が書かれていますので、戻る際にはその色の情報を-1してあげます。

こんな感じで、枝分かれしたところまで戻ったら、再び最下層を目指して探索をしていきます。

以上の探索を全ての頂点を通るまで行えば答えが求まります。

最後に本手法のアルゴリズムを図に示します。

実装

 #include <bits/stdc++.h> #define  rep(i,n) for(int i=0;i<n;++i) #define  reps(i,s,n) for(int i=s;i<n;++i)
using namespace std;
using ll = long long;
using P = pair<int, int>;

const int N = 1e5;
int c[N],color[N],used[N];
vector<vector<int>> g;
vector<int> ans;


void dfs(int a)
{
   int ca = c[a];
   //掲示板に書かれてなければ良い頂点
   //記録する
   if(color[ca] == 0) ans.emplace_back(a+1);
   used[a] = 1;
   //+1して次の頂点
   ++color[ca];
   for(auto t:g[a])
   {
       if(used[t] == 1) continue;
       dfs(t);
   }
   //戻るときには-1する
   --color[ca];
}

int main()
{
   int n;
   cin >> n;
   g.resize(n);
   rep(i,n) cin >> c[i];
   rep(i,n) --c[i];
   rep(i,n-1)
   {
       int a,b;
       cin >> a >> b;
       --a,--b;
       g[a].emplace_back(b);
       g[b].emplace_back(a);
   }
   
   dfs(0);
   sort(ans.begin(),ans.end());
   for(auto a:ans) cout << a << endl;

   return 0;
}

あとがき

掲示板を用意したDFSで解きました。

DFSは再帰関数を用いて実装しています。再帰関数の実装は慣れるまでに時間がかかるかもしれません。私がそうでした。難しいですもん。

ただ、本問題は戻り値を持たないで、外部の掲示板を更新しながら探索します。ですので、個人的には再帰関数の概念を考えやすい問題なのかと思います。もし再帰関数わからん！という方は少し時間をかけて本問題を考えてみると良い練習になるかと思います。なんか、理解できるときが急にきます。あれこうじゃないかな？という天啓が降りてきます。

再帰関数がわからなくても、whileとstackでもDFSは実装できますので、そちらでも良いかもしれません。まずは解けることが大切です！