【Python】平衡２分探索木の動作を Tkinter で可視化

# -*- coding:utf-8 -*-
import tkinter
import time
import math
import random

# キャンバスのサイズ
CANVAS_WIDTH = 1200
CANVAS_HEIGHT = 700

# データの最大数
DATA_NUM = 127

# 描画時のウェイト（s）
WAIT = 0.1

# 円の半径
RADIUS = 30


class Node():

	def __init__(self, number):
		self.left = None
		self.right = None
		self.number = number


class Tree():

	# 枝の方向を指す定数
	LEFT = 0
	RIGHT = 1

	# 処理を表す定数
	ADD = 0
	LEFT_ROTATE = 1
	RIGHT_ROTATE = 2
	DELETE = 3

	def __init__(self, view):
		'木構造を制御するオブジェクト生成'

		self.view = view

		# 木の根をNoneに初期化
		self.root = None

	def add_node(self, number):
		'ノードを追加'

		# まだノードが一つもない場合
		if not self.root:

			# 根ノードとしてノードを追加
			self.root = Node(number)
			self.draw(Tree.ADD)

			return True

		# 根ノードから順に追加する場所を探索
		node = self.root

		# 根ノードからどう辿ったかを覚えておくリスト
		branch = []
   
		while True:
			if number < node.number:
				# 追加する値がノードの値よりも小さい場合

				# 左の枝を辿ったことを覚えておく
				branch.append(Tree.LEFT)

				if not node.left:
					# そのノードの左の子が無い場合（もう辿るべきノードが無い場合）*/

					# その左の子の位置にノードを追加
					node.left = Node(number)

					# 追加完了したので処理終了
					break

				# 左の子がある場合は左の子を新たな注目ノードに設定
				node = node.left

			elif number > node.number:
				# 追加する値がノードの値よりも大きい場合

				# 右の枝を辿ったことを覚えておく
				branch.append(Tree.RIGHT)

				if not node.right:
					# そのノードの右の子が無い場合（もう辿るべきノードが無い場合）

					# その右の子の位置にノードを追加
					node.right = Node(number)

					# 追加完了したので処理終了
					break

				# 右の子がある場合は右の子を新たな注目ノードに設定
				node = node.right
			else:
				# 追加する値とノードの値が同じ場合

				print(str(number) + " already exist")
				return False

		# 一旦追加した状態で描画
		self.draw(Tree.ADD)

		# 平衡に保つための回転
		self.balancing(self.root, None, None, branch)

		return True

	def delete_node(self, number):
		'numberの値を持つノードを削除'

		if not self.root:
			return False

		# 削除するノードを探索
		node = self.root

		# 親ノードはひとまずNone
		parent = None

		# 根ノードからどう辿ったかを覚えておくリスト
		branch = []
		
		while node:
			if number < node.number:
				# 探索する値よりもnodeの値が大きい場合

				# 左の枝を辿ったことを覚えておく
				branch.append(Tree.LEFT)

				# 左の子ノードを辿る
				parent = node
				node = node.left
				  
			elif(number > node.number):
				# 探索する値よりもnodeの値が小さい場合

				# 右の枝を辿ったことを覚えておく
				branch.append(Tree.RIGHT)

				# 右の子ノードを辿る
				parent = node
				node = node.right
			else:
				# 見つかった場合
	  			break

		else:
			# 見つからなかった場合
			print(str(number) + "を持つノードはありません")
			return False
			
		if not node.left and not node.right:    
			# 子がいないノードの削除
			self.delete_no_child_node(node, parent)
		elif node.left and not node.right:
			# 左の子のみいるノードの削除
			self.delete_one_child_node(node, node.left)
		elif not node.left and node.right:
			# 右の子のみいるノードの削除
			self.delete_one_child_node(node, node.right)
		else:
			# 左の子と右の子両方がいるノードの削除
			self.delete_two_child_node(node, branch)

		# キャンバスから削除したノードの情報を除去
		self.remove(number)

		# 除去後の木を描画
		self.draw(Tree.DELETE)

		# 平衡にするため回転
		self.balancing(self.root, None, None, branch)

		return True

	def delete_no_child_node(self, node, parent):
		'子のいないノード(node)を削除'

		if parent:
			# 親がいる場合

			# nodeへの参照を外す
			if parent.left == node:
				parent.left = None
			else:
	  			parent.right = None
 
		else:
			# 親がいない場合
			self.root = None


	def delete_one_child_node(self, node, child):
		'子が一方にだけいるノード(node)の削除'

		# childeのデータをnodeにコピー
		node.number = child.number
		node.left = child.left
		node.right = child.right


	def delete_two_child_node(self, node, branch):
		'子が２つあるノード(node)の削除'

		# nodeの左の子孫ノードの中から最大値を持つノードを探索
		max_node = node.left
		max_parent = node

		# 左の枝を辿ったことを覚えておく
		branch.append(Tree.LEFT)

		# 必ず右の子ノードの方が値が大きい
		while max_node.right:
			max_parent = max_node
			max_node = max_node.right

			# 右の枝を辿ったことを覚えておく
			branch.append(Tree.RIGHT)

		# 最大ノードのデータ(number)をnodeにコピー
		node.number = max_node.number

  		# 最大ノード(max_node)を削除
  
		# 最大ノードなので右の子ノードはいない前提
		if not max_node.left:
			# max_nodeに子がいない場合
			self.delete_no_child_node(max_node, max_parent)   
		else:
			# 子が１つある場合
			self.delete_one_child_node(max_node, max_node.left)

	def balancing(self, node, parent, direction, branch):
		'各ノードの平衡係数を±1以下にする'

		if not node:
			return

		if len(branch) > 0:
			# 辿れる場合はまず目的のノードまで辿る

			# 辿る子ノードを設定
			if branch[0] == Tree.LEFT:
				next = node.left
			else:
				next = node.right

			# 子ノードを辿る
			self.balancing(next, node, branch[0], branch[1:])

		# 平衡係数を計算
		left_height = self.get_height(node.left)
		right_height = self.get_height(node.right)
		balance = right_height - left_height
		
		if balance > 1:
			# 右の部分木が高くて並行状態でない場合

			# ２重回転が必要かどうかを判断
			if self.get_height(node.right.left) > self.get_height(node.right.right):
				# ２重回転（Right Left Case）
				self.right_left_rotate(node, parent, direction)

			else:
				# １重回転（左回転）
				self.left_rotate(node, parent, direction)
			
		
		elif balance < -1:
			# 左の部分木が高くて並行状態でない場合

			# ２重回転が必要かどうかを判断
			if self.get_height(node.left.right) > self.get_height(node.left.left):
				# ２重回転（Left Right Case）
				self.left_right_rotate(node, parent, direction)
			else:
				# １重回転（右回転）
				self.right_rotate(node, parent, direction)
			

	def right_left_rotate(self, node, parent, direction):
		'２重回転（Right Left Case）を行う'

		print("right_left_rotate:" + str(node.number))

		# nodeの右の子ノードを根として右回転
		self.right_rotate(node.right, node, Tree.RIGHT)

		# nodeを根として左回転
		self.left_rotate(node, parent, direction)

	def left_right_rotate(self, node, parent, direction):
		'２重回転（Left Right Case）を行う'

		print("left_right_rotate:" + str(node.number))

		# nodeの左の子ノードを根として左回転
		self.left_rotate(node.left, node, Tree.LEFT)

		# nodeを根として右回転
		self.right_rotate(node, parent, direction)


	def left_rotate(self, node, parent, direction):
		'nodeを根として左回転を行う'

		print("left_rotate:" + str(node.number))

		# 新しい根とするノードをpivotとして設定
		pivot = node.right

		# 左回転
		if pivot:
			node.right = pivot.left
			pivot.left = node
		
		# parentもしくはrootに新しい根ノードを参照させる
		if not parent:
			self.root = pivot
			self.draw(Tree.LEFT_ROTATE)
			return

		# どちらの子に設定するかはdirectionから判断
		if direction == Tree.LEFT:
			parent.left = pivot
		else:
			parent.right = pivot

		self.draw(Tree.LEFT_ROTATE)

	def right_rotate(self, node, parent, direction):
		'nodeを根として右回転を行う'

		print("right_rotate:" + str(node.number))

		# 新しい根とするノードをpivotとして設定
		pivot = node.left

		# 右回転
		if pivot:
			node.left = pivot.right
			pivot.right = node

		# parentもしくはrootに新しい根ノードを参照させる
		if not parent:
			self.root = pivot
			self.draw(Tree.RIGHT_ROTATE)
			return

		# どちらの子に設定するかはdirectionから判断
		if direction == Tree.LEFT:
			parent.left = pivot
		else:
			parent.right = pivot

		self.draw(Tree.RIGHT_ROTATE)

	def get_height(self, node):
		'nodeを根とした木の高さを取得する'

		if not node:
			# nodeが無いなら高さは0
			return 0
		
		# 左右の子を根とした木の高さを取得
		left_height = self.get_height(node.left)
		right_height = self.get_height(node.right)

		# 大きい方に+1したものを木の高さとして返却
		if left_height > right_height:
			tree_height = left_height
		else:
			tree_height = right_height

		return tree_height + 1

	def draw(self, process):
		'木の描画'

		# rootを根とする木を描画
		self.draw_node(self.root, [], process)

		# キャンバスのアップデート依頼
		self.view.update()

		# スリープ
		time.sleep(WAIT)
		

	def draw_node(self, node, branch, process):
		'枝を辿りながらnodeを根とする木を描画'
		if not node:
			return

		# 処理によって円の色を変える
		if process == Tree.LEFT_ROTATE:
			diff_color = "#FFDDDD"
		elif process == Tree.RIGHT_ROTATE:
			diff_color = "#DDDDFF"
		elif process == Tree.ADD:
			diff_color = "#DDFFDD"
		elif process == Tree.DELETE:
			diff_color = "#FFFFDD"

		# そのノードと親ノードを繋ぐ線を描画
		self.view.draw_line(node, branch)

		# そのノードを描画
		self.view.draw_oval(node, branch, diff_color)

		# そのノードの数字を描画
		self.view.draw_text(node, branch)

		# 左の子孫ノードを描画
		branch.append(0)
		self.draw_node(node.left, branch, process)
		del branch[-1]

  		# 右の子孫ノードを描画
		branch.append(1)
		self.draw_node(node.right, branch, process)
		del branch[-1]

	def remove(self, number):
		'numberをデータに持つノードの描画を除去'

		self.view.remove_oval(number)
		self.view.remove_line(number)
		self.view.remove_text(number)

class View():

	def __init__(self, master):
		'UI関連のオブジェクト生成'

		self.master = master

		# 各種設定
		self.drawn_oval = []
		self.drawn_line = []
		self.drawn_text = []

		# キャンバスのサイズを決定
		self.canvas_width = CANVAS_WIDTH
		self.canvas_height = CANVAS_HEIGHT

		# 情報表示用のフレームを作成
		self.canvas_frame = tkinter.Frame(
			master,
		)
		self.canvas_frame.grid(column=1, row=1)

		# 操作用ウィジェットのフレームを作成
		self.operation_frame = tkinter.Frame(
			master,
		)
		self.operation_frame.grid(column=2, row=1, padx=10)

		# キャンバスの生成と配置
		self.canvas = tkinter.Canvas(
			self.canvas_frame,
			width=self.canvas_width,
			height=self.canvas_height,
			bg="#F0F0F0"
		)
		self.canvas.pack()

		# 開始ボタンの生成と配置
		self.button = tkinter.Button(
			self.operation_frame,
			text="開始",
		)
		self.button.pack()

	def get_oval_coord(self, branch):
		'キャンバス上の円の中心座標を取得'

		# まずは木のどの位置の円を描画するかをbranchから計算

		# branchの長さよりこのノードの縦方向の位置を計算
		tree_y = len(branch)

		# branchの要素より左から何番目のデータかを計算
		tree_x = 0
		i = 0
		for n in branch[::-1]:
			tree_x += n * 2 ** i
			i += 1


		# tree_xとtree_yをキャンバス上の座標に変換していく

		# データ数より木の最大高さを計算
		max_tree_height = math.ceil(math.log2(DATA_NUM + 0.5)) + 1

		# マスの数を計算
		num_horizontal = 2 ** (max_tree_height - 1) * 2
		num_vertical = max_tree_height * 2

		# １マスの幅を計算（キャンバスの幅とマス数より）
		one_width = self.canvas.winfo_width() / num_horizontal

		# １マスの高さを計算（キャンバスの高さとマス数より）
		one_height = self.canvas.winfo_height() / num_vertical


		# 円の中心座標を計算
		x_space = num_horizontal / (2 ** tree_y * 2)
		x = x_space * (tree_x * 2 + 1) * one_width
		y = (tree_y * 2 + 1) * one_height

		return x, y

	def judge_same_coord(self, coord1, coord2):
		'２つの座標が同じである角かを判断'

		if not coord1 or not coord2:
			# どちらか一方が None ならFalse
			return False

		if len(coord1) != len(coord2):
			# 長さが違うならFalse
			return False
		
		# 座標が同じかを１つ１つチェック
		for i in range(len(coord1)):
			if not math.isclose(coord1[i], coord2[i]):
				# １つでも異なればFasle
				return False

		# 全部同じならTrue
		return True


	def draw_oval(self, node, branch, diff_color):
		'ノードを表す円をdiff_color色で描画'

		# 円の中心座標を取得
		x, y = self.get_oval_coord(branch)

		# 円の左上（x1, y1）と円の右下（x2, y2）を計算
		coord = (
			x - RADIUS,
			y - RADIUS,
			x + RADIUS,
			y + RADIUS,
		)

		# タグを作成
		tag = "oval_" + str(node.number)

		# 前回描画時の円の座標を取得
		bcoord = self.canvas.coords(tag)

		# 前回と座標差分のあるノードだけ色を付ける
		if self.judge_same_coord(coord, bcoord):
			color = "#FFFFFF"
		else:
			color = diff_color
		
		# 前回描画した円を削除
		if tag in self.drawn_oval:
			self.drawn_oval.remove(tag)
			self.canvas.delete(tag)

		# 円を新しく描画
		self.canvas.create_oval(
			coord[0], coord[1],
			coord[2], coord[3],
			fill=color,
			tag=tag
		)

		# 描画した円のタグを覚えておく
		self.drawn_oval.append(tag)
		
	
	def draw_line(self, node, branch):
		'ノードと親ノードを繋ぐ線を描画'

		if len(branch) == 0:
			# タグを作成
			tag = "line_" + str(node.number)
		
			# 前回描画時した線を削除
			if tag in self.drawn_line:
				self.drawn_line.remove(tag)
				self.canvas.delete(tag)
			
			# 親がないので線は描画しない
			return

		# ノードの中心座標取得
		x, y = self.get_oval_coord(branch)

		# 親ノードの中心座標取得
		px, py = self.get_oval_coord(branch[:-1])
		
		width = px - x
		height = py - y
		rad = math.atan2(height, width)
		
		# タグを作成
		tag = "line_" + str(node.number)
		
		# 前回描画時した線を削除
		if tag in self.drawn_line:
			self.drawn_line.remove(tag)
			self.canvas.delete(tag)

		# 新しく線を描画
		self.canvas.create_line(
			x + RADIUS * math.cos(rad),
			y + RADIUS * math.sin(rad),
			px + RADIUS * math.cos(math.pi + rad),
			py + RADIUS * math.sin(math.pi + rad),
			tag=tag
		)

		# 描画した線のタグを覚えておく
		self.drawn_line.append(tag)

	def draw_text(self, node, branch):
		'ノードの数字を描画'

		# ノードの中心座標取得
		x, y = self.get_oval_coord(branch)

		# タグを作成
		tag = "text_" + str(node.number)
		
		# 前回描画時した数字を削除
		if tag in self.drawn_text:
			self.drawn_text.remove(tag)
			self.canvas.delete(tag)

		# 新しくテキストを描画
		self.canvas.create_text(
			x, y,
			text=str(node.number),
			tag=tag,
			font=("", RADIUS, "bold")
		)

		# 描画したテキストのタグを覚えておく
		self.drawn_text.append(tag)

	def remove_oval(self, number):
		'numberを持つノードの円を除去'

		# タグを作成
		tag = "oval_" + str(number)

		# 前回描画時した線を削除
		if tag in self.drawn_oval:
			self.drawn_oval.remove(tag)
			self.canvas.delete(tag)


	def remove_line(self, number):
		'numberを持つノードの線を除去'

		# タグを作成
		tag = "line_" + str(number)

		# 前回描画時した線を削除
		if tag in self.drawn_line:
			self.drawn_line.remove(tag)
			self.canvas.delete(tag)

	def remove_text(self, number):
		'numberを持つノードの数字を除去'

		# タグを作成
		tag = "text_" + str(number)

		# 前回描画時した円を削除
		if tag in self.drawn_text:
			self.drawn_text.remove(tag)
			self.canvas.delete(tag)

	def update(self):
		'キャンバスのアップデート'

		self.canvas.update()


class Controller():
	def __init__(self, view, tree):
		'SortとViewを制御するオブジェクトを生成'

		# 制御するViewとTreeのオブジェクト設定
		self.view = view
		self.tree = tree

		self.num_list = []

		# ボタンクリック時のイベントを受け付け
		self.view.button["command"] = self.button_click

	def button_click(self):
		'ボタンクリック時の処理'

		for i in range(DATA_NUM):
			self.add(i)

		for i in range(DATA_NUM):
			self.delete(i)


	def add(self, i):
		'木へのノードの追加を行う'

		num = random.randint(0, DATA_NUM)

		# 木へのnumを持つノード追加
		self.tree.add_node(num)
		print("add:" + str(num))

	def delete(self, i):
		'木へのノードの削除を行う'

		num = random.randint(0, DATA_NUM)

		# 木からnumを持つノードを削除
		self.tree.delete_node(num)
		print("delete:" + str(num))


# アプリ生成
app = tkinter.Tk()
app.title("二分木")

# Viewオブジェクト生成
view = View(app)

# Treeオブジェクト生成
tree = Tree(view)

# Controllerオブジェクト生成
controller = Controller(view, tree)

# mainloopでイベント受付を待機
app.mainloop()

# キャンバスのサイズ
CANVAS_WIDTH = 1200
CANVAS_HEIGHT = 700

# データの最大数
DATA_NUM = 127

# 描画時のウェイト（s）
WAIT = 0.1

# 円の半径
RADIUS = 30

# branchの長さよりこのノードの縦方向の位置を計算
tree_y = len(branch)

# branchの要素より左から何番目のデータかを計算
tree_x = 0
i = 0
for n in branch[::-1]:
	tree_x += n * 2 ** i
	i += 1

# マスの数を計算
num_horizontal = 2 ** (max_tree_height - 1) * 2
num_vertical = max_tree_height * 2

# データ数より木の最大高さを計算
max_tree_height = math.ceil(math.log2(DATA_NUM + 0.5)) + 1

# １マスの幅を計算（キャンバスの幅とマス数より）
one_width = self.canvas.winfo_width() / num_horizontal

# １マスの高さを計算（キャンバスの高さとマス数より）
one_height = self.canvas.winfo_height() / num_vertical

# 円の中心座標を計算
x_space = num_horizontal / (2 ** tree_y * 2)
x = x_space * (tree_x * 2 + 1) * one_width
y = (tree_y * 2 + 1) * one_height

x = x_space * (tree_x * 2 + 1) * one_width

# 円の左上（x1, y1）と円の右下（x2, y2）を計算
coord = (
	x - RADIUS,
	y - RADIUS,
	x + RADIUS,
	y + RADIUS,
)

# 円を新しく描画
self.canvas.create_oval(
	coord[0], coord[1],
	coord[2], coord[3],
	fill=color,
	tag=tag
)

# 描画した円のタグを覚えておく
self.drawn_oval.append(tag)

# 前回描画した円を削除
if tag in self.drawn_oval:
	self.drawn_oval.remove(tag)
	self.canvas.delete(tag)

# 前回描画時の円の座標を取得
bcoord = self.canvas.coords(tag)

# 前回と座標差分のあるノードだけ色を付ける
if self.judge_same_coord(coord, bcoord):
	color = "#FFFFFF"
else:
	color = diff_color

# 新しく線を描画
self.canvas.create_line(
	x + RADIUS * math.cos(rad),
	y + RADIUS * math.sin(rad),
	px + RADIUS * math.cos(math.pi + rad),
	py + RADIUS * math.sin(math.pi + rad),
	tag=tag
)

# 新しくテキストを描画
self.canvas.create_text(
	x, y,
	text=str(node.number),
	tag=tag,
	font=("", RADIUS, "bold")
)

# 新しく線を描画
self.canvas.create_line(
	x + RADIUS * math.cos(rad),
	y + RADIUS * math.sin(rad),
	px + RADIUS * math.cos(math.pi + rad),
	py + RADIUS * math.sin(math.pi + rad),
	tag=tag
)

# 前回描画時した線を削除
if tag in self.drawn_oval:
	self.drawn_oval.remove(tag)
	self.canvas.delete(tag)

if not math.isclose(coord1[i], coord2[i]):
	# １つでも異なればFasle
	return False

self.left = None
self.right = None
self.number = number