Tìm cặp phần tử tổng cho trước từ 2 cây BST khác nhau trong Python

Cho 2 cây BST và số Sum. Tìm các cặp (x,y) với x từ cây 1, y từ cây 2 sao cho x + y = Sum.

Ví dụ: Sum=12, cây 1 và cây 2 cho trước → [(6,6), (7,5), (9,3)].

Giải thuật (Duyệt in-order + 2 con trỏ)

• Duyệt in-order 2 cây BST → 2 mảng đã sắp xếp trav1, trav2.
• Hai con trỏ: left=0 trên trav1, right=len(trav2)-1 trên trav2.
• Nếu tổng = Sum → lưu cặp, left++, right--.
• Nếu tổng < Sum → left++.
• Nếu tổng > Sum → right--.

Ví dụ

class ListNode:
    def __init__(self, data):
        self.data = data
        self.left = self.right = None

def insert(root, key):
    if not root:
        return ListNode(key)
    if root.data > key:
        root.left = insert(root.left, key)
    else:
        root.right = insert(root.right, key)
    return root

def storeInorder(ptr, traversal):
    if not ptr:
        return
    storeInorder(ptr.left, traversal)
    traversal.append(ptr.data)
    storeInorder(ptr.right, traversal)

def solve(trav1, trav2, Sum):
    left, right = 0, len(trav2) - 1
    res = []
    while left < len(trav1) and right >= 0:
        s = trav1[left] + trav2[right]
        if s == Sum:
            res.append((trav1[left], trav2[right]))
            left += 1
            right -= 1
        elif s < Sum:
            left += 1
        else:
            right -= 1
    return res

def get_pair_sum(root1, root2, Sum):
    trav1, trav2 = [], []
    storeInorder(root1, trav1)
    storeInorder(root2, trav2)
    return solve(trav1, trav2, Sum)

root1 = None
for x in [9,11,4,7,2,6,15,14]:
    root1 = insert(root1, x)
root2 = None
for x in [6,19,3,2,4,5]:
    root2 = insert(root2, x)
print(get_pair_sum(root1, root2, 12))

Đầu vào

[9,11,4,7,2,6,15,14], [6,19,3,2,4,5], 12

Đầu ra

[(6, 6), (7, 5), (9, 3)]