How can I get array output instead of memory location when printing TreeNode?

Question

I am looking at solutions of LeetCode problem 108. Convert Sorted Array to Binary Search Tree:

Given an integer array nums where the elements are sorted in ascending order, convert it to a height-balanced binary search tree.

I'm trying to execute the solution code in my local IDE environment of Jupyter or PhyCharm. But it prints the result differently... How should I correct the code? Or is there's anything I didn't configure accurately?

#class TreeNode:
#    def __init__(self, val=0, left=None, right=None):
#        self.val = val
#        self.left = left
#        self.right = right


class Solution:
    def sortedArrayToBST(self, nums):
        if not nums:
            return None
        mid = len(nums) // 2
        root = TreeNode(nums[mid])
        root.left = sortedArrayToBST(nums[:mid])
        root.right = sortedArrayToBST(nums[mid+1:])
        
        return root

#case1
nums1 = [-10, -3, 0, 5, 9]
s = Solution()
root1 = s.sortedArrayToBST(nums1)
print(root1) # <__main__.TreeNode object at 0x000002BEA77FCE10>

I would like that last print call to output [0,-3,9,-10,null,5]

Another example:

#case2
nums2 = [1, 3]
root2 = s.sortedArrayToBST(nums2)
print(root2) # I would like it to print [3,1]

How can I achieve the tree to output in that format?

Answer 1

LeetCode serializes the value returned by your function in JSON format. If you want something similar to happen when you print a TreeNode outside of that LeetCode platform, then you'll have to write that serialization code yourself.

One way to do this is to extend the TreeNode class with a __repr__ method, and perform a BFS traversal on the tree that is rooted by the current node. Finally use the json library to convert the result to JSON:

Here is the adaptation to make to the TreeNode class:

import json
from collections import deque

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

    def __repr__(self):
        size = 0
        lst = []
        queue = deque([self])
        while queue:
            node = queue.popleft()
            if not node:
                lst.append(None)
            else:
                lst.append(node.val)
                size = len(lst)
                queue.append(node.left)
                queue.append(node.right)
    
        return json.dumps(lst[:size])

With this change, your print calls will have the desired output.