SQLite3で木構造の再現

以前javaで生成した二分探索技のデータをpythonで可視化しましたが、この流れを全てpythonでやってみます。この際、ノードの関係性をデータベースで再現します。

BinaryTreeクラスの作成

走査するアルゴリズムはJavaの時と同じです。



class Node:
    def __init__(self,data) -> None:
        self.data = data
        self.left = None
        self.right = None
        pass

    def getData(self):
        return self.data

    def setLeft(self,leftNode):
        self.left = Node(leftNode)
    
    def getLeft(self):
        return self.left
    
    def setRight(self,rightNode):
        self.right = Node(rightNode)

    def getRight(self):
        return self.right

class BinaryTree:
    def __init__(self,data) -> None:
        self.root = Node(data)
        self.btSTR = ""
    
    def addAuto(self,data) -> None:
        currNode = self.root
        while True:
            if data <= currNode.getData():
                if currNode.getRight() != None:
                    currNode = currNode.getRight()
                else:
                    currNode.setRight(data)
                    break
            else:
                if currNode.getLeft() != None:
                    currNode = currNode.getLeft()
                else:
                    currNode.setLeft(data)
                    break

    def out(self,node):
        if node.getLeft() != None:
            self.btSTR += f"{node.getData()}->{node.getLeft().getData()}\n"
            self.out(node.getLeft())
        if node.getRight() != None:
            self.btSTR += f"{node.getData()}->{node.getRight().getData()}\n"
            self.out(node.getRight())

データベースの作成

今回はルートノードをプライマリーキーとする、LeftNode,RightNodeの二つのカラムを持ったテーブルを作成します。

from BTgen import BinaryTree
import sqlite3

def main():
    arr = [19,21,22,25,16,13,7,34,24,5,31,9,2,48]
    bt = BinaryTree(20)
    for i in arr:
        bt.addAuto(i)
    bt.out(bt.root)
    return bt.btSTR[:-1]

dbname = "BT.db"
db = sqlite3.connect(f"./{dbname}")
c = db.cursor()

table_name = 'binarytree'

try:
    c.execute(
        f"""CREATE TABLE {table_name}(
            Root INTEGER,
            Left INTEGER,
            Right INTEGER
            )"""
    )
except sqlite3.OperationalError:
    pass

print(main())

for i in main().split("\n"):
    _r = i.split("->")
    root = int(_r[0])
    child = int(_r[1])

    if child< root:
        Sql =  f'INSERT INTO {table_name}("Root","Left","Right") VALUES({root},{child},NULL)'
    else:
        Sql =  f'INSERT INTO {table_name}("Root","Left","Right") VALUES({root},NULL,{child})'
    c.execute(
          Sql
    )


db.commit()
c.close()
db.close()

上記の登録方法では、以下のように被りが出てしまうため、もう少しコンパクトにできそうです。

データの追加を工夫します。

raw = main()
rows = raw.split("\n")

for i in set([int(i)for j in rows  for i in j.split("->")]):
    c.execute( f'INSERT INTO {table_name}("Root","Left","Right") VALUES({i},NULL,NULL)')

for i in rows:
    _r = i.split("->")
    root = int(_r[0])
    child = int(_r[1])

    if child< root:
        Sql =  f'UPDATE {table_name} SET "Left" = {child} WHERE  "root" = {root}'
    else:
        Sql =  f'UPDATE {table_name} SET "Right" = {child} WHERE  "root" = {root}'

    c.execute(
          Sql
    )

SQL文の生成を上記のように書き換えると、コンパクトになりました。

データベースからグラフの生成

登録したデータを使用して、グラフを作成します。

edges = {}
n = 0
for i in c.execute(f'SELECT * FROM {table_name}'):
    root = i[0]
    if i[1] != None:
        edges[(root,i[1])] = n
        n+=1
    if i[2] != None:
        edges[(root,i[2])] = n
        n+=1

G = Digraph(format="png")
G.attr("node", shape="circle")

for i,j in edges:
    G.edge(str(i), str(j))
G.render("bt")

edgesをSELECT文を生成して抽出するだけです。