AtCoder Library Practice Contest 'A - Disjoint Set Union'時間超過

Question

sampleは通るのですが、提出するとTLEになります。

def csearch(z, w):
	m = 0
	i = 0
	while i < len(con):
		if z in con[i]:
			if w in con[i]:	
				m += 1
			break
		else:
			 i += 1
	return m

def search(a):
	h = 0
	while h < len(con):
		if a in con[h]:
			break
		else:
			h += 1
	return h	
		
def add_edge(x, y):
	if search(x) < len(con):
		if search(y) < len(con):
			Y = search(y)
			con[search(x)].extend(con[Y])
			del con[Y]
		else:
			con[search(x)].append(y)
	else:
		if search(y) < len(con):
			con[search(y)].append(x)
		else:
			con.append([x, y])

N, Q = map(int, input().split())

con = [[0]]

for i in range(Q):
	t, u, v=map(int, input().split())
	
	if t == 0:
		if csearch(u, v) == 0:
			add_edge(u, v)
		
	elif t == 1:
		print(csearch(u, v))

Answer 1

TLEというのは時間超過のことなのですね。
ざっとみた感じでは、必要のないことをたくさんやっていることが原因でしょう。
より高速にするためにはデータの持ち方を変えないと難しいと思います。
私がやるなら、完全な仕分けはせずに回答に必要なところだけを行えるようなデータ構造にするでしょう。

どういうデータの持ち方が良いのかを説明してしまうと、Watarungurunnnさんの楽しみを奪うことになるので、ここに書くのはやめておきます。(人によってデータをどう持つのかはいろいろあるでしょうから)

Comments

[Watarungurunnn]

配慮のあるご回答ありがとうございます。
「データの持ち方」というのはつまりlistによるデータ表現を変えて検索をしやすくするべきという理解でよろしいでしょうか？

Answer 2

興を削ぐようで申し訳ないのですが、この問題にはすでに確立した解法が存在します。
必要であれば、「union find」で検索してみてください。