D – Minimum Steiner Tree
頂点に 1 から N の番号がついた N 頂点の木が与えられます。i 番目の辺は頂点 A_i と頂点 B_i を結んでいます。
このグラフからいくつかの(0 個でもよい)辺と頂点を削除してできる木のうち、指定された K 個の頂点、頂点
V_1, …, V_K を全て含むようなものの頂点数の最小値を求めてください。入力されるデータ
N K
A_1 B_1
A_2 B_2
：
A_N-1 B_N-1
V_1 V_2 … V_K（ただし、1 ≦ N, K ≦ 2×10^5
1 ≦ A_i, B_i ≦ N
1 ≦ V_1 ＜ V_2 ＜ … ＜ V_K ≦ N ）
解法はズバリ、木の葉（隣接する頂点がひとつしかない頂点）であって V に含まれない頂点とそれに接続する辺を削除することを可能な限り繰り返した結果できる木が求める木です。
ただし頂点数は最大で2×10^5もあります。また葉を取り除くことで葉ではなかった部分が新たに葉になる場合もあります。そこで木から葉を取り除く処理を高速で繰り返すコードを書かなければなりません。
Edgeクラスの定義
最初にEdgeクラスを定義します。木の頂点は無向辺でつながっていますが、ここでは双方向の有効辺にします。どの頂点からどの頂点に向かうのかと自分と反対向きの辺をすぐに取得できるようにフィールド変数に格納します。
|
1 2 3 4 5 6 7 8 9 10 11 |
class Edge { public Edge(int from, int next) { From = from; Next = next; } public int From = 0; public int Next = 0; public Edge Reverse = null; } |
Treeクラスの定義
Treeクラスを定義します。このクラスを実装することで、辺を追加する処理、頂点が葉なのかどうかを確認する、葉を取り除く処理ができるようになります。
フィールド変数は以下のとおりです。辺をすばやく削除できるように通常のリストではなくLinkedListを使います。普通のListでも試してみましたが TLE（制限時間オーバー） してしまうケースがあります。
|
1 2 3 4 5 6 |
class Tree { int N = 0; LinkedList<Edge>[] EdgeLists = new LinkedList<Edge>[0]; Dictionary<Edge, LinkedListNode<Edge>> NodesDic = new Dictionary<Edge, LinkedListNode<Edge>>(); } |
コンストラクタを示します。引数は頂点の数です。引数は頂点の数が渡されたら LinkedList の配列を生成して LinkedList のインスタンスを生成して格納します。
|
1 2 3 4 5 6 7 8 9 10 |
class Tree { public Tree(int n) { N = n; EdgeLists = new LinkedList<Edge>[N]; for (int i = 0; i < N; i++) EdgeLists[i] = new LinkedList<Edge>(); } } |
辺を追加する処理を示します。
ふたつの頂点の番号が渡されたらEdgeオブジェクトをふたつ生成します。そしてお互いの参照をEdge.Reverseにセットします。これで逆向きの辺を相互に取得できるようになります。
そしてLinkedListにオブジェクトを追加するのですが、削除するときに指定するLinkedListNodeを辞書に格納します。普通のListでRemoveするときの計算量はO(N)ですが、LinkedListであればO(1)です。辞書からLinkedListNodeを取得するときの計算量O(logN)を考慮してもこちらのほうが高速です。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
class Tree { public void AddDirectedEdge(int a, int b) { Edge edge1 = new Edge(a, b); Edge edge2 = new Edge(b, a); edge1.Reverse = edge2; edge2.Reverse = edge1; LinkedListNode<Edge> node1 = EdgeLists[a].AddLast(edge1); LinkedListNode<Edge> node2 = EdgeLists[b].AddLast(edge2); NodesDic.Add(edge1, node1); NodesDic.Add(edge2, node2); } } |
指定された番号の頂点が葉であるかどうかを調べる処理を示します。これは出次数（頂点から出ている辺の数）が1であるかどうかを調べるだけでわかります。
|
1 2 3 4 5 6 7 |
class Tree { public bool IsLeaf(int num) { return EdgeLists[num].Count == 1; } } |
すべての葉を取得するのであればこれでできます。葉である頂点の番号のリストが返されます。ただ何度もすべての葉を取得する処理を繰り返すのは時間がかかるのでこの問題を解くときには使いません。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
class Tree { public List<int> GetLeafs() { List<int> leafs = new List<int>(); for (int i = 0; i < N; i++) { if (IsLeaf(i)) leafs.Add(i); } return leafs; } } |
葉を削除する処理を示します。指定された番号の頂点が葉であるか確認してから削除の処理をおこないます。
まずその頂点から出ている辺がどの頂点に向かっているかを調べます。この頂点は葉なのでここから出ている辺と辺が向かっている頂点（コード内ではnext）はひとつしかありません。
nextが取得できたら辺のReverseを求めてこれをEdgeLists[next]から削除します。このときに呼び出すRemoveメソッドに渡す引数は辞書から求めることができます。そのあとEdgeLists[num]内の要素をクリアします。これで指定された頂点から出ている辺と入ってくる辺を削除することができました。
葉を削除することでこれまで葉ではなかった頂点 next が葉になるかもしれません。IsLeafメソッドでこれを調べて新たに葉になる頂点が存在するのであればその頂点番号を返します。そうでない場合は -1 を返します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
class Tree { public int RemoveLeaf(int num) { if (!IsLeaf(num)) return -1; int next = EdgeLists[num].First().Next; LinkedListNode<Edge> node = NodesDic[EdgeLists[num].First().Reverse]; EdgeLists[next].Remove(node); EdgeLists[num].Clear(); if(IsLeaf(next)) return next; else return -1; } } |
問題を解く
定義したクラスを用いて問題を解きます。
辺を追加したあと、残すK個の頂点番号が入力されますが、頂点番号 i が配列 V のなかに存在するかを素早く確認できるようにexistV[i]を参照すればわかるようにしています。
すべての葉を取得してこれを削除、またすべての葉を取得してこれを削除という処理を繰り返すのは明らかに効率が悪いです。そこでQueueのなかに 0 ～ N – 1 の値で配列 V 内に存在しないものを追加し、ひとつずつ値を取り出して葉かどうかを確認して葉であれば削除します。葉を削除したときに新たに葉ができたときはRemoveLeafメソッドの戻り値を見ればわかるので、これが配列Vに含まれない場合はQueueに追加します。そしてQueueが空になるまで処理を繰り返します。
問題の答えは配列Vに含まれない葉とこれに接続する辺を削除することを可能な限り繰り返してできる木の頂点数です。そこで変数 ans を N で初期化します。最初は頂点はひとつも削除されていないので ans = N です。頂点が削除されるたびにansをデクリメントします。ループから抜けたときのansの値が出力すべき答えです。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
class Program { static void Main() { int N, K; { int[] vs = Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); N = vs[0]; K = vs[1]; } Tree tree = new Tree(N); // 辺を追加 for (int i = 0; i < N - 1; i++) { // 入力は 1-indexed なので 0-indexed に変換するために -1 する int[] vs = Console.ReadLine().Split().Select(_ => int.Parse(_) - 1).ToArray(); int a = vs[0]; int b = vs[1]; tree.AddDirectedEdge(a, b); } // 入力は 1-indexed なので 0-indexed に変換するために -1 する int[] V = Console.ReadLine().Split().Select(_ => int.Parse(_) - 1).ToArray(); // 値が V に含まれているかすぐにわかるようにする bool[] existV = new bool[N]; for (int i = 0; i < V.Length; i++) existV[V[i]] = true; Queue<int> q = new Queue<int>(); // V に含まれていないすべての頂点番号をQueueにEnqueueする for (int i = 0; i < N; i++) { if (!existV[i]) q.Enqueue(i); } int ans = N; while (q.Count > 0) { int cur = q.Dequeue(); // Dequeueされた頂点番号が葉なら取り除く if (tree.IsLeaf(cur)) { int nextLeaf = tree.RemoveLeaf(cur); ans--; // 新たに取り除くべき葉が生まれたのであればEnqueueする if (nextLeaf != -1 && !existV[nextLeaf]) q.Enqueue(nextLeaf); } } Console.WriteLine(ans); } } |
