[C#] 第一回 アルゴリズム実技検定 K – 巨大企業

[C#] 第一回 アルゴリズム実技検定 K – 巨大企業 を解きたい

C# で参加した第一回 アルゴリズム実技検定(PAST)は、AからI問題まで解けました。第二回も参加するので予習がてら解いてない過去問を解いてみます。今回は第一回のK問題を解きます。

[C#] AtCoder 第一回 アルゴリズム実技検定 の振り返り │ Web備忘録

[C#] J – 地ならし [第一回 アルゴリズム実技検定] │ Web備忘録

AからJまで上の記事で解いてみたコードがあります。

K – 巨大企業

K – 巨大企業

ある企業には N 人の社員、社員 1,…,N が在籍する。このうち一人は社長であり、上司を持たない。その他の社員はみなちょうど一人の直属の上司を持つ。あなたのデータベースによると、社員 i (1≦i≦N) の直属の上司は社員 pi である (ただし社員 i が社長である場合は pi=−1 となっている)。この上下関係に循環は存在しない。

このデータを用いて、二つの社員番号 a,b を入力すると社員 a が社員 b の部下であるかを判定するサービスを作りたい。ここで、社員 a が社員 b の部下であるとは、社員 a から直属の上司を辿っていくことで社員 b に到達できることを意味する。

サービスに入力される Q 個の社員番号の組 (a1,b1),…,(aQ,bQ) が与えられる。これらのそれぞれに対して Yes または No を出力するプログラムを作成せよ。

制約

• 2≦N≦150,000
• 各 i=1,…,N に対し、pi=−1 または 1≤pi≤Npi=−1 なる i はちょうど 1 つ存在する。
• 社員間の上下関係に循環は存在しない。
• 1≦Q≦100,000
• 1≦ai,bi≦N
• ai≠bi

グラフ問題です。循環のない連結なグラフ、根付き木が与えられます。Q個のクエリが投げられ2つの頂点が親子関係にあるかどうかを高速に判定しなければなりません。

制約はN, Qともに10⁵なので単純な実装でクエリのたびに探索すると、O(N²) になってTLEします。

考えたこと

最終的に私が考えた解法は解説の実装方針と微妙に異なりますが以下のようなものです。

• クエリの内容を先読みしておく。
• 根から行きがけ順に探索し、現在探索している頂点の親をすべて配列に記録する。
• 探索中の頂点を子とするクエリがあれば、その親があるか判定して記録しておく。

行きがけ順に探索していくというところは解説の内容とあってはいます。がクエリを先に読んでおくので汎用性は低い実装かもしれません。

自力での実装

提出 #11606299 – 第一回 アルゴリズム実技検定 過去問

static void Main(string[] args)
{
    // 親から子に向けて辺を記録しておく
    N = sc.ReadInt();
    var root = 0;
    G = new List<int>[N];
    for (int i = 0; i < N; i++)
    {
        G[i] = new List<int>();
    }
    for (int i = 0; i < N; i++)
    {
        var a = sc.ReadInt() - 1;
        if (a < 0)
        {
            // 根の判定
            root = i;
            continue;
        }
        G[a].Add(i);
    }

    // クエリを先読みしておく
    var Q = sc.ReadInt();
    // クエリの入力を0-indexedのタプルで持っておく
    var QAB = sc.ReadTupArray(Q, -1);
    Query = new List<int>[N];
    for (int i = 0; i < N; i++)
    {
        Query[i] = new List<int>();
    }
    for (int i = 0; i < Q; i++)
    {
        // クエリを子ごとに記録しておく
        var c = QAB[i].Item1;
        var p = QAB[i].Item2;
        Query[c].Add(p);
    }

    // 最終的な答えを記録するセット
    Yes = new HashSet<Tuple<int, int>>();

    // 行きがけ順に親を記録しながら探索する
    Parents = new bool[N];
    Rec(root);

    // 記録された答えを参照して出力
    foreach (var cp in QAB)
    {
        if (Yes.Contains(cp))
        {
            Console.WriteLine("Yes");
        }
        else
        {
            Console.WriteLine("No");
        }
    }
}
static int N;
static List<int>[] G;
static HashSet<Tuple<int, int>> Yes;
static bool[] Parents;
static List<int>[] Query;
static void Rec(int now)
{
    // 今見ている頂点を子とするクエリをすべて走査
    foreach (var p in Query[now])
    {
        // 今見ている頂点の親を確認してクエリにこたえておく
        if (Parents[p]) Yes.Add(Tuple.Create(now, p));
    }
    // 今見ている頂点を親とする
    Parents[now] = true;
    foreach (var next in G[now])
    {
        // 行きがけ順にすべての子を参照する
        Rec(next);
    }
    // すべての子を探索し終えたら親のフラグを消す
    Parents[now] = false;
}

行きがけ順にすべての親を管理できるというのが大きな気づきでした。

1─┬─2─┬─3
  │   └─4
  └─5

行きがけ順に根付き木を探索すると上のような順になると思います。私の実装では訪問時に必要なクエリにこたえるというものです。

例えば3や4に来た時には 1 2 だけが親であり、2の探索を終えた段階では2の親フラグを消すので5に移ったときは1だけが親と判定されます。こんな感じで各頂点の親を訪問時にO(1)で判定できるようにしています。したがって訪問時に必要なクエリすべてにこたえてしまえばいいやと思った次第です。

たぶん上記実装で、計算量が O(N+Q) だと思います。

解説動画を参考にした実装

解説動画の実装方針は「行きがけ順で頂点を列挙すると、ある頂点の部分木に当たる頂点が連続した区間に並ぶ」ということを利用すればクエリに O(1) でこたえられるというものです。

1─┬─2─┬─3
  │   └─4
  └─5

たしかに連続した区間に部分木のノードが並びます。ということでスマートな実装かどうかはわからないですが、各頂点の部分木のサイズと、列挙したリスト上でのインデックスがわかれば判定できそうです。

ということで解説動画を見た後に実装したのがこちらです。

提出 #11606032 – 第一回 アルゴリズム実技検定 過去問

static void Main(string[] args)
{
    // 頂点から子の向きに辺を張る
    N = sc.ReadInt();
    G = new List<int>[N];
    for (int i = 0; i < N; i++)
    {
        G[i] = new List<int>();
    }
    var root = 0;
    for (int i = 0; i < N; i++)
    {
        var p = sc.ReadInt() - 1;
        if (p < 0)
        {
            root = i;
            continue;
        }
        G[p].Add(i);
    }

    // 頂点から部分木を列挙したリスト
    PartialTree = new List<int>();
    // ある頂点を根とした部分木のサイズ
    PartialTreeSize = new int[N];
    // 列挙したリストにおける頂点のインデックス
    PartialTreeIndex = new int[N];

    // 根から部分木を列挙する
    Rec(root);

    var list = new List<string>();
    var Q = sc.ReadInt();
    for (int i = 0; i < Q; i++)
    {
        // 部下の頂点が上司の頂点を根とする部分木に含まれるかどうかを判定する
        var c = sc.ReadInt() - 1;
        var p = sc.ReadInt() - 1;
        // リストのインデックスと部分木のサイズをもとに区間を見る
        var l = PartialTreeIndex[p];
        var r = l + PartialTreeSize[p] - 1;
        var ci = PartialTreeIndex[c];
        if (l <= ci && ci <= r)
        {
            list.Add("Yes");
        }
        else
        {
            list.Add("No");
        }
    }
    foreach (var ans in list)
    {
        Console.WriteLine(ans);
    }
}
static int N;
static List<int>[] G;
static List<int> PartialTree = new List<int>();
static int[] PartialTreeSize;
static int[] PartialTreeIndex;
static int Rec(int now)
{
    // 現在の頂点が部分木のリストでインデックスを記録
    PartialTreeIndex[now] = PartialTree.Count;
    // 部分木のリストに頂点を追加
    PartialTree.Add(now);

    // 再帰的に部分木を走査してサイズを計算しておく
    var size = 1;
    foreach (var next in G[now])
    {
        size += Rec(next);
    }

    // 現在の頂点の部分木のサイズを記録して返す
    PartialTreeSize[now] = size;
    return size;
}

こっちのほうが実行速度も速く、クエリにも柔軟に対応できてよさげですね。

以上。