Python's nonlocal keyword - is this good practice?

Question

Consider a simple situation such as finding the k'th smallest element in a BST.

In my solution below:

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        i = 0
        ans = -1
        def traverse_inorder(root):
            nonlocal i
            nonlocal ans
            if not root:
                return
            traverse_inorder(root.left)
            i += 1
            if i == k:
                ans = root.val
                return
            traverse_inorder(root.right)
        traverse_inorder(root)
        return ans

Is my use of nonlocal for i and ans good practice? I do this to keep track of how many elements I've traversed after reaching the left most node (smallest value).

Another solution would be to have i and ans as member variables of the class:

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.i = 0
        self.ans = -1
        def traverse_inorder(root):
            # etc etc

Are both methods equivalent? Is one practice better than the other, and why?

Answer 1

I think that under the circumstances you describe, the nonlocal approach objectively makes more sene than the other one you show. Your goal is to have a counter outside your function and a way to keep track of the result when it is found regardless of where you are in the recursion.

nonlocal completely encapsulates that information in a namespace dedicated to the particular run of your outer function. There is really no downside here.

Using instance attributes makes that information available to anyone using the instance. This is unnecessary conceptually, marginally slower, and not thread-safe. While thread safety is often not a concern in Python, it does make your code much less robust. Between that and encapsulation, I would definitely stay away from this approach.

I would suggest a third possibility here: using return values. In that case, your really don't need any external namespace. I don't necessarily recommend this over using nonlocal, but it's worth mentioning. Here is a sample implementation, which as you can see is much more verbose than your solution:

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def traverse_inorder(root, target):
            if not root:
                return 0, None
            left_count, item = traverse_inorder(root.left, target)
            if left_count == target - 1:
                 return left_count + 1, root.val
            elif left_count < target:
                 right_count, item = traverse_inorder(root.right, target - left_count - 1)
                 return left_count + right_count + 1, item
            else:  # left_count == target
                 return left_count, item
        count, ans = traverse_inorder(root, k)
        if count < k:
            raise ValueError('Insufficient elements')
        return ans

There are many ways to do this with return values. Here, I compute the number of elements in each subtree, up to the target value. A non-none item is only returned if the exact number of elements is found.