二分木探索をpythonで作って学ぼう！

2024年6月22日

こんにちは！

この記事では二分木探索について解説します。G検定などDeep learningの勉強をしていると必ず出てくるのが二分木探索です。

幅優先探索、深さ優先探索。。

これらを解説文だけ読んで理解するのは難しいですよね。なので今回はこれらの探索をpythonで視覚的に表現してみます！
文字で見るよりわかりやすいと思うので興味のある方は続きを読んでみてください！

作成したコードも載せます！解説はしません！

幅優先探索と深さ優先探索

二分探索木を使った探索にはいくつか種類があります。代表的なのが幅優先探索と深さ優先探索です。

これらは探索のアルゴリズムが違います。

それぞれメリットとデメリットを押さえておくのが大事なので今回学んでいきましょう！

それぞれ最初にプログラムを実行したデモ動画を載せます！どんな順番で探索をしているのかを意識しながら見てください。

幅優先探索

これが幅優先探索です。

開始地点から全ての隣接ノードを順番に探索していきます。

今回の例でいうと、①が開始地点です。①の隣接ノードは②と③なのでこれらを探索します。隣接ノードを全て探索したら次は②の隣接ノードの探索、③の隣接ノードの探索というように進みます。

特徴としては、探索の順番を書いたときに数字が横に昇順で並ぶことです。

次に、幅優先探索のメリットとデメリットを紹介します。

幅優先探索のメリット

最短経路を見つけることができる
実装が簡単

まず幅優先探索は近いノードから順番に探索をしていくので目的のノードまでの最短経路を確実に見つけることができます。

また開発に目を向けたときに、キューを利用するだけでいいのでアルゴリズムの実装が容易です。

幅優先探索のデメリット

メモリ消費量が多い
遅い

まず幅優先探索では目的ノードを見つけるまでのノードを全てキューに保持しておく必要があり、メモリの消費量が多くなります。

また全ノードを探索するため、目的地点に到達するまでの時間も長くなります。

深さ優先探索

これが深さ優先探索です。開始地点の隣接ノードを順番に一番深くまで探索していきます。

先ほどの幅優先探索の動画と見比べると違いがよくわかると思います。深さ優先探索にもメリット・デメリットはあるので見ていきましょう。

深さ優先探索のメリット

メモリ効率が良い
全ての経路を探索

深さ優先探索のメリットはメモリ効率が良いことです。

幅優先探索ではノードの深さごとに探索するため、ゴールが見つかるまでの経路を全て保持しておく必要がありメモリ消費量が多くなります。

一方深さ優先探索では、ノードごとに一番深くまで探索するためその経路に目的地がなければその情報は削除してしまって構わないため、現在探索している経路の情報だけをメモリに保持しておけば良いのです。
そのためメモリ効率が非常によくなります。

二つ目のメリットに「全ての経路を探索」と書きましたが、これもメモリ効率の良さからきています。

とても深いノードに目的地がある場合を考えてみてください。幅優先探索の場合、全経路を保持しておく必要があるため目的地が深くにあればあるほどメモリの使用量が増えてしまいます。

その結果メモリが足りなくて探索が異常終了してしまうこともあります。

対して、深さ優先探索は探索中の経路だけを保持しておけば良いためより多くの経路を探索することができます。

大事なのは深さ優先探索はメモリ効率が良いということです。

深さ優先探索のデメリット

最短経路が保証されない

これは説明しなくても明らかですが、深さごとに探索をする訳ではないので最短経路を必ず見つけることができるわけではないです。

幅優先探索と深さ優先探索

まとめると、

point!

幅優先探索はメモリ効率は悪いが最短経路を見つけることができる。

深さ優先探索はメモリ効率はいいが最短経路は保証されない。

ということです！

ここだけは覚えておきましょう。この後プログラムも公開します。ちなみにですがこのプログラムを実行するとそれぞれのメモリ使用量も見ることができるようにしてあります。
気になる方はチェックしてみてください！

プログラム公開

では今回作成したプログラムのソースコードを公開します。言語はpythonです。

import pygame
import sys
import tracemalloc
from collections import deque

# ノードの定義
class Node:
    def __init__(self, key=None, left=None, right=None):
        self.key = key
        self.left = left
        self.right = right

# 二分木を描画する関数
def draw_tree(node, x, y, dx, dy, surface, font):
    if node:
        # ノードの円を描画
        pygame.draw.circle(surface, (node.color), (x, y), 20)
        # ノードの値を描画
        if node.key is not None:
            text = font.render(str(node.key), True, (0, 0, 0))
            surface.blit(text, (x-10, y-10))
        # 左子ノードが存在する場合、線を描画し再帰的に呼び出し
        if node.left:
            pygame.draw.line(surface, (255, 255, 255), (x, y), (x-dx, y+dy), 1)
            draw_tree(node.left, x-dx, y+dy, dx//2, dy, surface, font)
        # 右子ノードが存在する場合、線を描画し再帰的に呼び出し
        if node.right:
            pygame.draw.line(surface, (255, 255, 255), (x, y), (x+dx, y+dy), 1)
            draw_tree(node.right, x+dx, y+dy, dx//2, dy, surface, font)

# 幅優先探索を行いノードを訪問する関数
def bfs(root, screen, size, font):
    tracemalloc.start()
    queue = deque([root])
    order = 1
    while queue:
        node = queue.popleft()
        node.key = order
        node.color = (0, 255, 0)  # 探索済みノードの色を緑に設定
        order += 1

        # 再描画
        screen.fill((0, 0, 0))
        draw_tree(root, size[0]//2, 50, 200, 100, screen, font)
        pygame.display.flip()
        pygame.time.delay(500)  # 500msの遅延

        if node.left:
            queue.append(node.left)
        if node.right:
            queue.append(node.right)
    
    current, peak = tracemalloc.get_traced_memory()
    print(f"BFS Memory usage: Current={current / 10**6}MB; Peak={peak / 10**6}MB")
    tracemalloc.stop()

# 深さ優先探索を行いノードを訪問する関数
def dfs(node, order, screen, size, font):
    if node:
        node.key = order[0]
        node.color = (0, 255, 0)  # 探索済みノードの色を緑に設定
        order[0] += 1

        # 再描画
        screen.fill((0, 0, 0))
        draw_tree(root, size[0]//2, 50, 200, 100, screen, font)
        pygame.display.flip()
        pygame.time.delay(500)  # 500msの遅延

        dfs(node.left, order, screen, size, font)
        dfs(node.right, order, screen, size, font)

# Pygameの初期化
pygame.init()
size = (800, 600)
screen = pygame.display.set_mode(size)
pygame.display.set_caption("Binary Tree")
font = pygame.font.SysFont(None, 24)

# 二分木の作成（もう一層追加）
root = Node(None,
            Node(None,
                 Node(None,
                      Node(), Node()),
                 Node(None,
                      Node(), Node())),
            Node(None,
                 Node(None,
                      Node(), Node()),
                 Node(None,
                      Node(), Node())))

# ノードの初期色を設定
def set_initial_color(node):
    if node:
        node.color = (255, 255, 255)  # 白色
        set_initial_color(node.left)
        set_initial_color(node.right)

set_initial_color(root)

# メインループ
running = True
started = False
algorithm = None

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_b:
                algorithm = "bfs"
                started = True
            elif event.key == pygame.K_d:
                algorithm = "dfs"
                started = True

    if started:
        if algorithm == "bfs":
            bfs(root, screen, size, font)
        elif algorithm == "dfs":
            tracemalloc.start()
            order = [1]  # 探索順序のカウンタ
            dfs(root, order, screen, size, font)
            current, peak = tracemalloc.get_traced_memory()
            print(f"DFS Memory usage: Current={current / 10**6}MB; Peak={peak / 10**6}MB")
            tracemalloc.stop()
        started = False  # 一度実行したら停止

    screen.fill((0, 0, 0))
    draw_tree(root, size[0]//2, 50, 200, 100, screen, font)
    pygame.display.flip()

pygame.quit()
sys.exit()