AVL树实现新冠感染数据结构：WorseBefore函数实现求助

Great job picking an AVL tree for this problem—it’s the right choice for maintaining ordered data with O(log n) insertions and queries. The challenge with WorseBefore(d) is narrowing down the largest date less than d where the sick count exceeds d's count, and your existing node structure can be extended to make this efficient.

The key issue with your current setup is that the maxLeftSick and maxRightSick don’t help you quickly rule out subtrees that can’t contain valid candidates, especially since those max values might include dates larger than d. We need to add a few more pieces of metadata to each node to enable pruning during the query.

Step 1: Extend Your AVL Node Structure #

Update each node to track three additional values:

• min_key: The smallest date in the subtree rooted at this node
• max_key: The largest date in the subtree rooted at this node
• max_sick: The highest sick count in the subtree rooted at this node

These values are easy to maintain during insertion and rotations. For any node, you can compute them as:

# Update min_key for the node
node.min_key = node.date
if node.left:
    node.min_key = min(node.min_key, node.left.min_key)
if node.right:
    node.min_key = min(node.min_key, node.right.min_key)

# Update max_key for the node
node.max_key = node.date
if node.left:
    node.max_key = max(node.max_key, node.left.max_key)
if node.right:
    node.max_key = max(node.max_key, node.right.max_key)

# Update max_sick for the node
node.max_sick = node.sick
if node.left:
    node.max_sick = max(node.max_sick, node.left.max_sick)
if node.right:
    node.max_sick = max(node.max_sick, node.right.max_sick)

After inserting a new node or performing a rotation, traverse up from the modified node to the root, updating these three values for each ancestor. This ensures all metadata stays consistent.

Step 2: Implement the WorseBefore(d) Function #

First, retrieve the target sick count by finding the node with date d (assume d exists per the problem’s examples—handle missing d as needed). Then, use a recursive search to find the largest valid date:

Recursive Implementation

def find_max_valid(node, target_date, target_sick):
    if not node:
        return None
    
    if node.date >= target_date:
        # Only check left subtree if it could have valid candidates
        if node.left and node.left.max_sick > target_sick:
            return find_max_valid(node.left, target_date, target_sick)
        else:
            return None
    else:
        # First check right subtree for larger valid dates
        right_result = None
        if node.right and node.right.min_key < target_date and node.right.max_sick > target_sick:
            right_result = find_max_valid(node.right, target_date, target_sick)
        
        if right_result is not None:
            return right_result
        
        # Check current node if right subtree has no valid candidates
        if node.sick > target_sick:
            return node.date
        
        # Finally check left subtree for smaller valid dates
        if node.left and node.left.max_sick > target_sick:
            return find_max_valid(node.left, target_date, target_sick)
        else:
            return None

def WorseBefore(d):
    # Get the target sick count for date d
    target_node = avl_tree.search(d)
    if not target_node:
        return None  # Adjust behavior per problem requirements
    target_sick = target_node.sick
    return find_max_valid(avl_tree.root, d, target_sick)

How It Works

• When the current node’s date is >= d, we only check the left subtree (since all right subtree dates are larger). We skip the left subtree entirely if its max sick count is <= target_sick—no valid candidates exist there.
• When the current node’s date is < d, we first check the right subtree (to find larger dates than the current node that are still < d). We only recurse into the right subtree if it contains dates < d and has a max sick count > target_sick.
• If the right subtree has no valid candidates, we check if the current node itself is valid.
• If not, we check the left subtree (for smaller valid dates, our last resort if no larger valid dates exist).

Step 3: Verify with Your Example #

Let’s walk through your first WorseBefore(13) call:

1. Target sick count is 18 (from date 13).
2. The search starts at the root. When it reaches node 5 (date <13), it checks the right subtree (node 8). Node 8’s right subtree (date 13) has min_key >=13, so no recursion there. Node 8’s sick count (13) is <18, so we check its left subtree (null).
3. Back to node 5, we check its left subtree (node 2). Node 2’s right subtree (node3) has max_sick=15 <=18, so no recursion. Node2’s sick count (20) is >18, so we return 2.

After inserting date10 (sick=19), the second WorseBefore(13) call will find node10 in the left subtree of node8, which has a sick count >18 and date <13—so it returns 10 as expected.

Why This Is O(log n) #

Each recursive call traverses one level of the AVL tree, and we prune entire subtrees whenever we know they can’t contain valid candidates. Since the AVL tree’s height is O(log n), the query runs in O(log n) time.

内容的提问来源于stack exchange，提问作者user14826913