TLEが数ケースだけ発生していた話

ContainsでTLEが発生してた話

ABC368のD問題、意気揚々と解いたが２件だけTLE。

ABC368 D

            int[] NK = Console.ReadLine().Split(' ').Select(int.Parse).ToArray();
            SortedList<int, List<int>> rootTree = new SortedList<int, List<int>>();
            int[] score = new int[NK[0]];

            if (NK[1] == 1)
            {
                // 残す要素数1の時、1が確定
                Console.WriteLine(1);
                return;
            }
            else if (NK[0] == NK[1])
            {
                // 残す要素数 == 全体の要素数の場合もN==Kが確定する
                Console.WriteLine(NK[0]);
                return;
            }

            // 初期化
            for (int i = 0; i < NK[0]; i++)
            {
                rootTree[i] = new List<int>();
            }

            // 辺の読み込み
            for (int i = 0; i < NK[0] - 1; i++)
            {
                int[] ab = Console.ReadLine().Split(' ').Select(int.Parse).ToArray();
                var a = ab[0] - 1;
                var b = ab[1] - 1;
                rootTree[a].Add(b);
                rootTree[b].Add(a);
            }

            // 残す要素を読み込む
            int[] vk = Console.ReadLine().Split(' ').Select(s => int.Parse(s) - 1).ToArray();

            // 作成した根付き木に対し、スコアを計算する
            Func<int, int, int>? pruning = null;
            pruning = (myself, parent) =>
            {
                foreach (var child in rootTree[myself])
                {
                    if (child != parent)
                    {
                        score[myself] += pruning!(child, myself);
                    }
                }
                if (vk.Contains(myself))
                {
                    score[myself]++;
                }
                return score[myself];
            };

            pruning(vk[0], -1);

            Console.WriteLine(score.Where(w => w > 0).Count());

で、一時間くらい悩んでほかの人の解法まで見に行ってようやく気が付いた。
Containsを使って判定しているところ、最大だと計算量O（N）かかるじゃん。

            pruning = (myself, parent) =>
            {
                <<中略>>
                if (vk.Contains(myself))
                // このContainsはvk.Countが大きいほどO(N)に近づく。
                {
                    score[myself]++;
                }
                return score[myself];
            };

各項目で再帰しているため、全体の計算量は最大だとO(N^2)かかる。
で、たぶんTLEになった２ケースはこの最大に近い計算量がかかるケースになってたようだ。
（VK内の要素数がNの最大に近いんだと思う）

ここを以下のようにあらかじめ計算量O(1)で取得できるように変更するだけで無事ACとなった。

            int[] vk = Console.ReadLine().Split(' ').Select(s => int.Parse(s) - 1).ToArray();
+            bool[] vks = new bool[NK[0]];
+            foreach (int i in vk)
+            {
+                vks[i] = true;
+            }

            // 作成した根付き木に対し、スコアを計算する
            Action<int, int>? pruning = null;
            pruning = (myself, parent) =>
            {
-                if (vk.Contains(myself))
+                if (vks[myself])
                {
                    score[myself]++;
                }
                foreach (var child in rootTree[myself])
                {
                    if (child != parent)
                    {
                        pruning(child, myself);
                        score[myself] += score[child];
                    }
                }
            };

結論

Containsを使うときは、操作コレクションの要素数が大きくなるかどうかを検討してから使う必要がある。