https://qiita.com/bluejenga/items/3ee278abbf8702a5e728#comment-383873e2ef6804892708

tree.py

#!/usr/bin/env python3

from enum import IntEnum

import random
import itertools
import sys


# 隣のセルへの座標の差分 (x, y), [North, East, South, West]
DELTA = [(0, -1), (1, 0), (0, 1), (-1, 0)]


class DIR(IntEnum):
    N = 0  # North
    E = 1  # East
    S = 2  # South
    W = 3  # West

    def opposite(self):
        return (self.S, self.W, self.N, self.E)[self]


class Cell:

    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.clear()

    def clear(self):
        # 迷路の一部に割り当てられたかどうか
        self.used = False
        # 点灯してるかどうか
        self.light = False
        # 隣接するセルと接続している方角 [N, E, S, W]
        self._linked = [False, False, False, False]

    # def get__linked_str(self): => str(cell)
    def __str__(self):
        """接続している方角を文字で取得する"""
        n, e, s, w = self._linked
        return "N"*n + "E"*e + "S"*s + "W"*w

    def __getitem__(self, direction):
        return self._linked[direction]

    def __setitem__(self, direction, link):
        self._linked[direction] = link

    def __iter__(self):
        return cell._linked

    def link(self, direction, neighbor):
        self[direction] = True
        neighbor[direction.opposite()] = True
        neighbor.used = True

    def is_linked(self, direction, neighbor):
        return self[direction] and neighbor[direction.opposite()]

    def rotate(self, n=1):
        """セルを時計回りに90*n度回転する"""
        self._linked = self._linked[-n:] + self._linked[:-n]

    def shuffle(self):
        self.rotate(random.choice(range(len(DIR))))


class Tree:
    """
    木を表すクラス。
    長方形に並んだCellで構成される。
    不要な領域にはCellの代わりにNoneを割り当てて、木の形を表現する。
    CellとCellの接続が迷路の通路を表す。
    """

    def __init__(self, height):
        self.height = height
        # widthはheightから算出する。
        self.width = width = max(1, height * 2 - 3)
        self.center = center = width // 2

        self._cells = [None] * (width * height)
        for y, x in itertools.product(range(height), range(width)):
            self[x, y] = Cell(x, y)  # 無効な座標には設定されない

        # 根元、幹のCell
        self.root = self[center, height - 1]

    def index(self, x, y):
        """(x, y)座標を一次元配列インデックスに変換する"""
        return y * self.width + x

    # def get_cell_list(self): => for cell in tree:
    def __iter__(self):
        """有効なセルを列挙する"""
        return (cell for cell in self._cells if cell)

    def __getitem__(self, coord):
        if coord in self:
            x, y = coord
            return self._cells[self.index(x, y)]
        return None

    def __setitem__(self, coord, cell):
        if coord in self:
            x, y = coord
            self._cells[self.index(x, y)] = cell

    # def tree.is_valid_coord(self, x, y): => (x, y) in tree:
    def __contains__(self, coodinate):
        """指定した座標が存在するかどうかを返す"""
        x, y = coodinate
        # 根元
        if x == self.center and y == self.height - 1:
            return True

        # 根元より上の長方形の範囲に収まるかどうか
        if not (0 <= x < self.width and 0 <= y < self.height - 1):
            return False

        # 木の形(三角形)の範囲に収まるかどうか
        return abs(x - self.center) <= y

    def neighbor(self, cell, direction):
        """指定した方角の隣接するセルを取得する"""
        dx, dy = DELTA[direction.value]
        return self[cell.x + dx, cell.y + dy]

    def neighbors(self, cell, used=True):
        """cellセルを起点とする接続（通路）を列挙する"""
        return [(cell, direction, neighbor)
                for direction in DIR
                for neighbor in [self.neighbor(cell, direction)]
                if neighbor and neighbor.used == used]

    def build(self):
        """迷路を構築する"""
        # 初期化
        for cell in self:
            cell.clear()

        # 木の根元から迷路作成開始（実際にはどこからでもよい）
        self.root.used = True
        links = self.neighbors(self.root, used=False)

        # 通路候補がなくなるまで繰り返し
        while len(links) > 0:
            # 候補から通路を一つ選択して接続する
            src, direction, dst = random.choice(links)
            src.link(direction, dst)

            # 既存候補の中から、dstが重複しているものを削除
            links = [link for link in links if link[2] != dst]
            # dstを起点とした通路を候補に追加
            links.extend(self.neighbors(dst, used=False))

    def shuffle(self):
        """各セルの向きをシャッフルする"""
        for cell in self:
            cell != self.root and cell.shuffle()

    def rotate(self, x, y):
        """指定した座標のセルを時計回りに90度回転する"""
        cell = self[x, y]
        cell and cell.rotate()

    def lightup(self):
        """根元から辿って、繋がっているセルを点灯する"""
        def lightup_recursive(cell):
            cell.light = True
            # 通路が伸びている方角を取得
            for _, direction, neighbor in self.neighbors(cell):
                # そのセルからも自分の方に通路が伸びていれば、再帰的に点灯処理
                if cell.is_linked(direction, neighbor):
                    neighbor.light or lightup_recursive(neighbor)

        # 一旦すべて消灯してから、接続を辿って点灯
        for cell in self:
            cell.light = False
        lightup_recursive(self.root)

    def is_complete(self):
        """すべてのセルが点灯していればTrueを返す"""
        return all(cell.light for cell in self)

    def print(self):
        """コンソールに木の絵を出力"""
        def to_strings(cell):
            if not cell:
                return ['     '] * 3

            n = '   ' if cell[DIR.N] else '---'
            s = '   ' if cell[DIR.S] else '---'
            w = ' ' if cell[DIR.W] else '|'
            e = ' ' if cell[DIR.E] else '|'
            L = 'X' if cell.light else ' '
            return [f'+{n}+', f'{w} {L} {e}', f'+{s}+']

        print(' ', *[f'  {x:2}' for x in range(self.width)])

        cells = [to_strings(cell) for cell in self._cells]
        xs = range(self.width)
        for y in range(self.height):
            for i, label in enumerate(('   ', f'{y:2} ', '   ')):
                print(label, *[cells[self.index(x, y)][i] for x in xs], sep='')


def main(height):
    print('Enter Commands.')
    print('    X Y  Rotate cell')
    print('    n    Start new game')
    print('    e    Exit')
    print()

    tree = Tree(height)
    while True:
        tree.build()
        tree.shuffle()
        tree.lightup()
        tree.print()

        while True:
            command = input('>>> ')
            if command == 'e':
                return
            if command == 'n':
                break
            try:
                x, y = map(int, command.split())
                tree.rotate(x, y)
            except ValueError:
                pass

            tree.lightup()
            tree.print()

            if tree.is_complete():
                print('Congratulations!!!!!')


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('usage: tree.py <height>')
        exit()
    main(int(sys.argv[1]))