Neunomizuの日記

俺だけが俺だけじゃない

# LeetCode Medium 116. Populating Next Right Pointers in Each Node

tags: leetcode

問題

イデア

完全二分木が与えられた時に,あるノードと同じ深さにある右のノードをポインタnextで指せという問題です.

使用するメモリは定数倍でないといけないらしいですが,再帰ならセーフらしいです.今回は再帰の分がセーフということは木の深さ分までならメモリを使ってもいいということでやってみました.(スタック領域のみと書いてあります)

解法

last_visitedという連想配列を用意して,ある深さにおいて最後に訪れたノードを記憶します.

左部分木からノードを辿っていくことで,常に最後に訪れたノードが左側のノードになります.そのた,えこのノードのnextを現在のノードにすれば,題意を満たします.

計算量

ノードの数を$N$,高さを$H$とします.

  • 時間計算量
    • $O(N)$
  • 空間計算量
    • $O(H)$

Python

"""
# Definition for a Node.
class Node:
    def __init__(self, val: int = 0, left: 'Node' = None, right: 'Node' = None, next: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right
        self.next = next
"""
from typing import Mapping

class Solution:
    def connect(self, root: 'Node') -> 'Node':
        self.dfs(root, {}, 0)
        return root
    
    def dfs(self, node: 'Node', last_visited: Mapping[int, 'Node'], depth: int) -> 'Node':
        if not node:
            return node
        if depth in last_visited: 
            last_visited[depth].next  = node
        last_visited[depth] = node
        for vertex in [node.left, node.right]: 
            self.dfs(vertex, last_visited, depth+1)