ABC168のD問題で解法が同じなのにTLEしてしまったわけ【Java】

今回はABC168の復習をしました。

先日のABC168のD問題で、解法が全く同じにもかかわらずTLEを出し続けて沼にはまりました。

それが悔しかったので、この原因を調べてみました。

ちなみにABC168の結果がこちらです。
developer-bpeldi2oerkd8.hatenablog.com

コンテスト中に出した答え

コンテスト中に出した答えはこちらです。(長いですけど)

ちゃんとBFSを使って解いているにも関わらず、TLEしました。

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.Scanner;

public class Main {
	//java11

	public static void main(String[] args) {

		Scanner sc = new Scanner(System.in);

		int N = sc.nextInt();
		int M = sc.nextInt();

		Graph g = new Graph(N);
		for(int i=0; i<M; i++) {
			int A = sc.nextInt();
			int B = sc.nextInt();

			g.addEdge(g.getNode(A), g.getNode(B));
			g.addEdge(g.getNode(B), g.getNode(A));
		}

		boolean[] visited = new boolean[N];
		Arrays.fill(visited, false);
		//BFS
		Queue<Node> queue = new ArrayDeque<>();
		queue.add(g.getNode(1));
		while(!queue.isEmpty()) {
			Node n = queue.poll();

			//訪問済み
			visited[n.value-1] = true;

			for(Edge edge : n.edges) {
				Node node = edge.to;

				if(!queue.contains(node) && !visited[node.value-1]) {
					queue.add(node);
					if(node.root != null) {
						System.out.println("No");
						return;
					}else {
						node.root = n;
					}
				}
			}
		}

		System.out.println("Yes");
		for(int i=2; i<=N; i++) {
			System.out.println(g.getNode(i).root.value);
		}

	}

	//一般的なグラフ
	public static class Graph{
		int n;
		Map<Integer, Node> nodes;	//(ノード番号, ノード)

		//1~Nまでをもつグラフ
		public Graph(int n) {
			this.n = n;
			nodes = new HashMap<>(n);

			for(int i=0; i<n; i++) {
				nodes.put(i+1, new Node());
			}
		}

		public Graph(int n, int[] values) {
			this.n = n;
			nodes = new HashMap<>(n);

			for(int i=0; i<n; i++) {
				nodes.put(values[i], new Node(values[i]));
			}
		}

		public Node getNode(int node_n) {
			//ノード番号を指定してノードを取得
			return nodes.get(node_n);
		}

		public void addEdge(Node from, Node to) {
			//重みが1のグラフの場合(重みが均等な場合にも場合によっては使える)
			from.edges.add(new Edge(to, 1));
		}

		public void addEdge(Node from, Node to, long weight) {
			//エッジの重みが異なる場合
			from.edges.add(new Edge(to, weight));
		}

		public void removeEdge(Node from, Node to) {
			//指定したノード間のエッジを削除
			from.edges.remove(new Edge(to, 1));
		}
	}

	public static int count = 0;
	//ノードを表すクラス
	public static class Node{

		int value;
		List<Edge> edges;
		boolean visited = false;
		Node root;

		public Node(){
			value = count+1;
			edges = new ArrayList<Edge>();
			count++;
		}

		public Node(int value) {
			this.value = value;
			edges = new ArrayList<Edge>();
		}

		@Override
		public boolean equals(Object obj) {
			if(obj instanceof Node) {
				Node n = (Node)obj;
				if(this.value == n.value)
					return true;
				else
					return false;
			}

			return false;
		}

		@Override
		public int hashCode() {
			return Objects.hash(this.value);
		}
	}

	//エッジを表すクラス
	public static class Edge{
		Node to;		//行き先
		long weight;		//重み

		public Edge(Node to, long weight) {
			this.to = to;
			this.weight = weight;
		}

		@Override
		public boolean equals(Object obj) {
			if(obj instanceof Edge) {
				Edge e = (Edge)obj;
				if(this.to.equals(e.to))
					return true;
				else
					return false;
			}

			return false;
		}

		@Override
		public int hashCode() {
			return Objects.hash(this.to);
		}
	}

}

原因の分析

TLEになった原因をJavaの実行時間を測るSystem.nanoTimeを使って地道に調べました。

すると、意外にも自作のライブラリのほうではなく、BFSの実装部分のほうに原因があるとわかってきました。
（自作のライブラリが原因だと思ってかなりの時間を使ってしまったのは内緒です。）

さらに絞り込んでいくと、どうやらif(!queue.contains(node) && !visited[node.value-1])で時間を消費していることがわかってきました。

どう考えても、!visited[node.value-1]には時間がかかるはずがないので、（配列だからO(1)でとってこれるはず）
原因は!queue.contains(node)にあると推測しました。

TLEになった原因

結局TLEになった原因は、queueに多くの要素が含まれるようになったとき、queue.containsに多くの時間がかかってしまうことでした。

queue.containsメソッドの内部処理はよく知りませんが、要素を1つずつ確認していくのであればO(N)の時間がかかってしまうので、ここで時間がかかっていたのだと思います。

対処法としては、queueに加えるときに、visited=trueとして、訪問済みにしてしまうことです。

こうすれば、queueに加えるときに訪問済みかどうか調べるだけでいいので、O(1)の時間しかかかりません。

修正したプログラムはこちらです。
（自作のライブラリの部分は同じなので省略しています。）

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import java.util.Scanner;

public class Main {
	//java11

	public static void main(String[] args) {

		Scanner sc = new Scanner(System.in);

		int N = sc.nextInt();
		int M = sc.nextInt();

		Graph g = new Graph(N);
		for(int i=0; i<M; i++) {
			int A = sc.nextInt();
			int B = sc.nextInt();

			g.addEdge(g.getNode(A), g.getNode(B));
			g.addEdge(g.getNode(B), g.getNode(A));
		}

		boolean[] visited = new boolean[N];
		Arrays.fill(visited, false);
		//BFS
		Queue<Node> queue = new ArrayDeque<>();
		queue.add(g.getNode(1));
		visited[0] = true;
		while(!queue.isEmpty()) {
			Node n = queue.poll();

			for(Edge edge : n.edges) {
				Node node = edge.to;

				if(!visited[node.value-1]) {
					queue.add(node);
                                        //queueに加えたときに訪問済みにしてしまう
					visited[node.value-1] = true;
					if(node.root != null) {
						System.out.println("No");
						return;
					}else {
						node.root = n;
					}
				}
			}
		}

		System.out.println("Yes");
		for(int i=2; i<=N; i++) {
			System.out.println(g.getNode(i).root.value);
		}

	}

さらに、これだと制限時間ギリギリなので、入出力を変えてみました。

入力にネットで拾ってきたFastScannerを使ったバージョンと、それに加え、出力をPrintWriterにしたバージョンの実行時間を計ってみた結果がこちらです。

上から、

FastScanner+PrintWriter (563ms)
FastScannerのみ (1099ms)
そのまま (1648ms)

です。

いや、すごい威力ですね。

今回のように、入出力が長い行になる場合は迷わずFastScanner, PrintWriterを使ったほうがよさそうです。

というわけで、ABC168のD問題の復習でした。

BFSは知っていたのに時間内にACできなかったのは悔しいですが、結果的に勉強になりました。

BFSのテンプレを今回の形に修正しておきます。

次こそは、４完できるように頑張ります。