Utility functions for use in LeetCode

create_binary_tree.py

from typing import List
from collections import deque

# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def create_tree(nodes: List) -> TreeNode:
    """Build a binary tree based on the array input from LeetCode.
    
    :param nodes: The array of binary tree node values in the LeetCode format.
    :returns: The root node of the tree.
    """
    if not nodes:
        return None
    root = TreeNode(val=nodes[0])
    node_queue = deque()
    node_queue.append(root)
    i = 1
    while node_queue and i < len(nodes):
        curr = node_queue.popleft()
        curr.left = TreeNode(val=nodes[i]) if nodes[i] else None
        if i + 1 >= len(nodes):
            break
        curr.right = TreeNode(val=nodes[i + 1]) if nodes[i + 1] else None
        if curr.left:
            node_queue.append(curr.left)
        if curr.right:
            node_queue.append(curr.right)
        i += 2
    return root

# Sample useage
tree_array = [1, 3, 2, 5, None, None, 9, 6, None, None, 7]
root = create_tree(tree_array)

# This builds this tree:
#            1
#          /   \
#         3     2
#        /       \
#       5         9
#      /           \
#     6             7

create_multilevel_linked_list.py

from typing import List


# Definition for a Node.
class Node:
    def __init__(self, val, prev, next, child):
        self.val = val
        self.prev = prev
        self.next = next
        self.child = child


def convert_to_multilevel_lst(node_lst: List) -> List[List]:
    """Convert a multilevel linked list from 1D to multilevel array of arrays.

    :param node_lst: The 1D array repr of a multilevel linked list, in LeetCode
                    style.
    :returns: The multilevel array repr of the multilevel linked list.
    """
    node_lst.append(None)
    multi_node_lst = []
    is_new_level = True
    for i, val in enumerate(node_lst):
        if is_new_level:
            temp = []
            is_new_level = False
        temp.append(val)
        if val is None and node_lst[i - 1] is not None:  # end of curr level
            is_new_level = True
            multi_node_lst.append(temp)
    return multi_node_lst


def create_doubly_linked_list(node_lst: List) -> Node:
    """Create doubly linked list from the given node_lst in LeetCode style.
    
    This requirement is seen in the question of "Flatten a Multilevel Doubly 
    Linked List".
    
    :param node_lst: The list of node values for the doubly linked list in
                    LeetCode style.
    """
    def print_multilevel_linked_list() -> None:
        """Print the multilevel linked list in the following format.

        1->2->3->None
           |
           4->5->None
              |
              6->None

        It works only for single digit values.
        """
        for i, root in enumerate(levels):
            num_nones = 0
            for j in range(i + 1):
                for val in multilevel_lst[j]:  # count number of leading Nones
                    if val is None:
                        num_nones += 1
                    else:
                        break
            print('   ' * num_nones, end='')
            curr = root
            while curr:
                print(f'{curr.val}->', end='')
                curr = curr.next
            print(None)
            print('   ' * num_nones, end='')
            curr = root
            while curr:
                if curr.child:
                    print('|  ', end='')
                else:
                    print('   ', end='')
                curr = curr.next
            print()

    # turn original node_lst into multi-layer node_lsts
    multilevel_lst = convert_to_multilevel_lst(node_lst)
    levels = []
    for single_level in multilevel_lst:  # create linked lists for each level
        idx = 0
        while not single_level[idx]:  # skip leading Nones
            idx += 1
        root = Node(single_level[idx], None, None, None)
        idx += 1
        pre = root
        while single_level[idx]:
            curr = Node(single_level[idx], pre, None, None)
            pre.next = curr
            pre = curr
            idx += 1
        levels.append(root)
    # link linked lists from adjacent levels together
    for i, root in enumerate(levels[1:], 1):
        num_nones = 0
        for val in multilevel_lst[i]:  # count the number of leading Nones
            if val is None:
                num_nones += 1
            else:
                break
        pre_lvl_node = levels[i - 1]  # root node of the previous level
        # move to the node that connects to the current level.
        for _ in range(num_nones):
            pre_lvl_node = pre_lvl_node.next
        pre_lvl_node.child = root

    print_multilevel_linked_list()  # print to check whether the build is good

    return levels[0] if levels else None

# Sample usage
node_lst = [1,2,3,4,5,6,None,None,None,7,8,9,10,None,None,11,12]
root = create_doubly_linked_list(node_lst)
# If we print the multilevel linked list while calling create_doubly_linked_list(),
# we will get this output indicating the shape of the created doubly linked
# list:
# 1->2->3->4->5->6->None
#       |           
#       7->8->9->10->None
#          |        
#          11->12->None

disjoint_union.py

class DJU:
    """This disjoint union supports union and find (with path compression).

    It is based on the original index of the given elements, thus there is no
    need to set up an additional mapping from the element value to its index.
    """

    def __init__(self, n: int):
        """Constructor. 
        
        Create a list, where the value of the ith element denotes the root node 
        index of the union the ith element belongs to.
        
        :param n: Total number of elements.
        """
        self.unions = list(range(n))

    def find(self, idx) -> int:
        """Find root for the idx-th element.
        
        Perform path compression along the way.

        :param idx: The index whose root is to be found.
        :return: The pos of the current node n and the root.
        """
        if self.unions[idx] != idx:
            self.unions[idx] = self.find(self.unions[idx])  # compression
        return self.unions[idx]

    def union(self, i, j) -> None:
        """Merge the union of the ith and jth element.

        :param i: Index of one of the node to be merged.
        :param j: Index of the other node to be merged.
        """
        self.unions[self.find(i)] = self.find(j)