トップページ -> AOJの解答例 -> ALDS1_8の解答例

ALDS1_8の解答例(Python)

二分探索木に関する問題です. A,B,Cしかできていない上にコードが相当酷いですが,一応載せておきます. アルゴリズムについて詳しく知りたい方は参考文献などをご覧ください.

1. ALDS1_8_A: Binary Search Tree I

二分探索木にキーを挿入するinsertと先行順巡回と中間順巡回を出力するprintを実装する問題です. 先行順巡回と中間順巡回はALDS1_7で作ってあるのでそれを流用しました.

# ALDS1_8_A
# 左 根 右 葉まで下る 左側の要素がなければ(根節点なので)追加 それ以外の点は下から戻ってきたときに追加する
def inOrder(root):
    x = tree[root]
    node_list = []
    node_set = set()
    while x!=None:
        flag = True
        # 左側の要素があるかを調べる
        if x.left != None:
            if x.left.key not in node_set:
                x = x.left
                flag = False
        else:
            # 左側の要素がない場合は追加してよい
            if x.key not in node_set:
                node_list.append(x.key)
                node_set.add(x.key)
        # 無ければ右側を調べる
        if flag:
            flag = True
            # 葉までたどり着いたら追加
            if x.right == None:
                x = x.parent
                if x.key not in node_set:
                    node_list.append(x.key)
                    node_set.add(x.key)
                flag = False
            if flag:
                # まだ調べていない場合は下る
                if x.right.key not in node_set:
                    x = x.right
                # 既に調べてある場合は上る
                else:
                    x = x.parent
                    if x != None:
                        if x.key not in node_set:
                            node_list.append(x.key)
                            node_set.add(x.key)
        
    print(" ",end="") # 問題に合わせて空白を入れただけ
    print(*node_list)
    
# 下るときに追加していけばいい 葉までたどり着いたら右側の要素が追加されていないところまで戻る
def preOrder(root):
    x = tree[root]
    node_list = [root]
    node_set = set([root])
    while x!=None:
        # print(node_list)
        flag = True
        # 左側の要素があるかを調べる
        if x.left != None:
            if x.left.key not in node_set:
                x = x.left
                node_list.append(x.key)
                node_set.add(x.key)
                flag = False
                
        # 無ければ右側を調べる
        if flag:
            flag = True
            # 葉までたどり着いたら戻る
            if x.right == None:
                x = x.parent
                flag = False
            if flag:
                # まだ調べていない場合は下る
                if x.right.key not in node_set:
                    x = x.right
                    node_list.append(x.key)
                    node_set.add(x.key)
                # 既に調べてある場合は上る
                else:
                    if x != None:
                        x = x.parent
          
    print(" ",end="") # 問題に合わせて空白を入れただけ
    print(*node_list)

class Node:
    def __init__(self,key):
        self.key = key
        self.right = None
        self.left = None
        self.parent = None
        
def insert_key(node,root):
    key = node.key
    if not tree.keys():
        tree[key] = node
        return key # 根を返す
    else:
        x = tree[root]
        while True:
            # xより小さい場合
            if key < x.key:
                # 左側の要素がないなら追加して終了
                if x.left == None:
                    x.left = node
                    node.parent = x
                    tree[node.key] = node
                    tree[x.key] = x
                    break
                else:
                    x = x.left
                    
            elif x.key < key:
                # 右側の要素がないなら追加して終了
                if x.right == None:
                    x.right = node
                    node.parent = x
                    tree[node.key] = node
                    tree[x.key] = x
                    break
                else:
                    x = x.right
                    
        return root
        
tree = {}
root = None

n = int(input())
for _ in range(n):
    command = input().split(" ")
    if command[0] == "print":
        inOrder(root)
        preOrder(root)
    elif command[0] == "insert":
        key = int(command[1])
        root = insert_key(Node(key),root)

2. ALDS1_8_B: Binary Search Tree II

Aに追加でキーが存在するかどうかを判断するfindを実装します. 以下を追加するだけでキーが見つけられます.

def find(key):
    if key in tree.keys():
        print("yes")
    else:
        print("no")
        
tree = {}
root = None

n = int(input())
for _ in range(n):
    command = input().split(" ")
    if command[0] == "print":
        inOrder(root)
        preOrder(root)
    elif command[0] == "insert":
        key = int(command[1])
        root = insert_key(Node(key),root)
    elif command[0] == "find":
        key = int(command[1])
        find(key)

3. ALDS1_8_C: Binary Search Tree III

Bに追加してキーを削除するdeleteを実装します. コードの書き方が悪すぎてdeleteの際にキーがすり替わったりして混迷を極めています... 削除する対象の持つ子が1つ以下なら簡単なのですが,2つのときは中間順巡回で"適切なキー"(次のキー)を探す必要があります. 中間順巡回の処理が別々に書かれていたり,コメントが悪かったりとあまりにも酷いですが,一応通ります.

# ALDS1_8_C 通ったけどひどい

# deleteで子を2つ持つキーを削除する際に"適切なキー"を見つける関数
def findDeleteKey(root,key):
    x = tree[root]
    node_list = []
    node_set = set()
    while x!=None:
        flag = True
        # 左側の要素があるかを調べる
        if x.left != None:
            if x.left.key not in node_set:
                x = x.left
                flag = False
        else:
            # 左側の要素がない場合は追加してよい
            if x.key not in node_set:
                if key in node_set:
                    return x.key
                node_list.append(x.key)
                node_set.add(x.key)
        # 無ければ右側を調べる
        if flag:
            flag = True
            # 葉までたどり着いたら追加
            if x.right == None:
                x = x.parent
                if x != None:
                    if x.key not in node_set:
                        if key in node_set:
                            return x.key
                        node_list.append(x.key)
                        node_set.add(x.key)
                flag = False
            if flag:
                # まだ調べていない場合は下る
                if x.right.key not in node_set:
                    x = x.right
                # 既に調べてある場合は上る
                else:
                    x = x.parent
                    if x != None:
                        if x.key not in node_set:
                            if key in node_set:
                                return x.key
                            node_list.append(x.key)
                            node_set.add(x.key)
                            
    return None

# 左 根 右 葉まで下る 左側の要素がなければ(根節点なので)追加 それ以外の点は下から戻ってきたときに追加する
def inOrder(root):
    x = tree[root]
    node_list = []
    node_set = set()
    while x!=None:
        flag = True
        # 左側の要素があるかを調べる
        if x.left != None:
            if x.left.key not in node_set:
                x = x.left
                flag = False
        else:
            # 左側の要素がない場合は追加してよい
            if x.key not in node_set:
                node_list.append(x.key)
                node_set.add(x.key)
        # 無ければ右側を調べる
        if flag:
            flag = True
            # 葉までたどり着いたら追加
            if x.right == None:
                x = x.parent
                if x != None:
                    if x.key not in node_set:
                        node_list.append(x.key)
                        node_set.add(x.key)
                flag = False
            if flag:
                # まだ調べていない場合は下る
                if x.right.key not in node_set:
                    x = x.right
                # 既に調べてある場合は上る
                else:
                    x = x.parent
                    if x != None:
                        if x.key not in node_set:
                            node_list.append(x.key)
                            node_set.add(x.key)
        
    print(" ",end="") # 問題に合わせて空白を入れただけ
    print(*node_list)
    
# 下るときに追加していけばいい 葉までたどり着いたら右側の要素が追加されていないところまで戻る
def preOrder(root):
    x = tree[root]
    node_list = [root]
    node_set = set([root])
    while x!=None:
        # print(node_list)
        flag = True
        # 左側の要素があるかを調べる
        if x.left != None:
            if x.left.key not in node_set:
                x = x.left
                node_list.append(x.key)
                node_set.add(x.key)
                flag = False
                
        # 無ければ右側を調べる
        if flag:
            flag = True
            # 葉までたどり着いたら戻る
            if x.right == None:
                x = x.parent
                flag = False
            if flag:
                # まだ調べていない場合は下る
                if x.right.key not in node_set:
                    x = x.right
                    node_list.append(x.key)
                    node_set.add(x.key)
                # 既に調べてある場合は上る
                else:
                    if x != None:
                        x = x.parent
          
    print(" ",end="") # 問題に合わせて空白を入れただけ
    print(*node_list)

class Node:
    def __init__(self,key):
        self.key = key
        self.right = None
        self.left = None
        self.parent = None
        
def insert_key(node,root):
    key = node.key
    if not tree.keys():
        tree[key] = node
        return key # 根を返す
    else:
        x = tree[root]
        while True:
            # xより小さい場合
            if key < x.key:
                # 左側の要素がないなら追加して終了
                if x.left == None:
                    x.left = node
                    node.parent = x
                    tree[node.key] = node
                    tree[x.key] = x
                    break
                else:
                    x = x.left
                    
            elif x.key < key:
                # 右側の要素がないなら追加して終了
                if x.right == None:
                    x.right = node
                    node.parent = x
                    tree[node.key] = node
                    tree[x.key] = x
                    break
                else:
                    x = x.right
                    
        return root
    
def find(key):
    if key in key_dict.keys():
        key = key_dict[key]
        
    if key in key_set:
        print("yes")
    else:
        print("no")
        
def delete(root,key):
    if key in key_set:
        key_set.remove(key)
    if key in key_dict.keys():
        key = key_dict[key]
    
    children = (tree[key].left!=None) + (tree[key].right!=None)
    # 子要素がない場合 keyの親から見てkeyを消せばいい
    if children == 0:
        # 親の右側の要素の場合
        if key > tree[key].parent.key:
            tree[key].parent.right = None
        # 親の左側の要素の場合
        else:
            tree[key].parent.left = None
    # 子要素が1つの場合 keyの親の子(key)をkeyの子とすり替える
    elif children == 1:
        if tree[key].left!=None:
            child = tree[key].left
        elif tree[key].right!=None:
            child = tree[key].right
            
        # 親の右側の要素の場合
        if key > tree[key].parent.key:
            tree[key].parent.right = child
            child.parent = tree[key].parent
        # 親の左側の要素の場合
        else:
            tree[key].parent.left = child
            child.parent = tree[key].parent
    # 子要素が2つの場合 自分を消して"適切なキー"を代わりに入れる
    else:
        swap_key = findDeleteKey(root,key)
        delete(root,swap_key)
        
        tree[key].key = swap_key
        key_dict[swap_key] = key
            
tree = {}
key_dict = {}
key_set = set()
root = None

n = int(input())
for _ in range(n):
    command = input().split(" ")
    if command[0] == "print":
        inOrder(root)
        preOrder(root)
    elif command[0] == "insert":
        key = int(command[1])
        root = insert_key(Node(key),root)
        key_set.add(key)
    elif command[0] == "find":
        key = int(command[1])
        find(key)
    elif command[0] == "delete":
        key = int(command[1])

4. ALDS1_8_D: Treap

Treapに関する問題です.

# ALDS1_8_D
# 編集中

<- 前へ戻る 【目次に戻る】 次へ進む ->