AtCoder Beginner Contest 198 E Unique Color を再帰を使わないで書いた

1. 概要

AtCoder Beginner Contest 198 E Unique Color
の解答について、再帰を使わないで書いてみました。

#2. 問題文

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

#3. 実装方法
自作クラスとして treeNode を定義し、treeNode.parent が親のノード番号、treeNode.children が子のノード番号になるように、最初に DFS でデータを整理しています。探索の終了判定のために、木構造の原点となる root ノード(つまり頂点0)の親のノード番号は -1 としました。

アルゴリズムのメインの部分は、公式解説(https://atcoder.jp/contests/abc198/editorial/540 )にある通りです。子ノードに進むたびに色の総和の配列(colorAccum)をインクリメントし、親ノードに戻るたびに colorAccum をデクリメントしています。currentNode の値がその時点で注目しているノードの番号となっています。

#4. 提出結果
https://atcoder.jp/contests/abc198/submissions/22286137

指標	パフォーマンス
実行時間	251 ms
メモリ	17700 KB

#5. 提出コード

ABC0411_E.cpp

#include <iostream>
#include<string>
#include <vector>
#include <queue>
#include <stack>
#include <iomanip>
#include <stdio.h>
#include <math.h>

using namespace std;

class treeNode{
public:
    int id;
    int depth;
    int parent;
    vector<int> connect;
    vector<int> children;
    int color;
    int nextChildIndex;
    bool isGood;
    
    treeNode(){
        id = 0;
        depth = -1;
        parent = -1;
        color = 0;
        nextChildIndex = 0;
        isGood = false;
    }

    int getNextChild(){
        int result;
        if(children.size() > nextChildIndex){
            result = children[nextChildIndex];
            nextChildIndex++;
        } else {
            result = -1;
        }
        return result;
    }
};

int main()
{
    
    int n;
    cin >> n;
    int a;
    int b;
    vector<treeNode> tn(n); // tn: tree node
    stack<int> nodeStack;

    for (int i = 0; i < n; i++)
    {
        cin >> tn[i].color;
    }

    for (int i = 0; i < n-1; i++)
    {
        cin >> a >> b;
        tn[a-1].connect.push_back(b-1);
        tn[b-1].connect.push_back(a-1);
    }

    tn[0].depth = 0;
    nodeStack.push(0);
    while(nodeStack.size() > 0){ // SoA か AoS か、実行時間に差はない
        int origin = nodeStack.top();
        nodeStack.pop();
        int nodes = tn[origin].connect.size();
        for (int i = 0; i < nodes; i++)
        {
            int target = tn[origin].connect[i];
            if(tn[target].depth<0){
                tn[target].parent = origin;
                tn[origin].children.push_back(target);
                tn[target].depth = tn[origin].depth + 1;
                nodeStack.push(target);
            }
        }
    }
    // ここからアルゴリズムのメイン
    tn[0].isGood = true;
    int currentNode = 0;
    vector<int> colorAccum(100010, 0);
    colorAccum[tn[0].color]++;
    while(currentNode >=  0){
        int nextNode = tn[currentNode].getNextChild();
        if(nextNode >= 0){ // 子ノードに進む場合
            if(colorAccum[tn[nextNode].color] == 0){
                tn[nextNode].isGood = true;
            }
            colorAccum[tn[nextNode].color]++;
        } else { // 親ノードに戻る場合
            nextNode = tn[currentNode].parent;
            colorAccum[tn[currentNode].color]--;
        }
        currentNode = nextNode;
    }

    for (int i = 0; i < n; i++)
    {
        if(tn[i].isGood){
            cout << (i+1) << endl;
        }
    }

    return 0;
}

#6. 感想・コメント
再帰を使わなくても単純なスタックだけで上手く解けないかと思って試行錯誤していたのですが、なかなか難しかったので木構造のクラスを自作してみることにしました。模範解答のように再帰だけでサクっと書くのと比べると冗長ですが、初心者的には分かりやすく書けた気もします。

再帰が無いので、スタックオーバーフローの可能性は無くなったかと思います。ただ、そもそも競技プログラミングで再帰をループで書かないとエラーになるような問題は出そうにないので、競プロに慣れている人ならば再帰で書いたほうがよいでしょう。

あくまで、木構造を勉強中の方のご参考までに。