给定元素与插入位置，求大规模数据下的最终数组（优化超时问题）

Let's break down this problem clearly: we have two arrays arr and index (both of length n), and we need to insert each arr[i] into the current array at position index[i]—shifting existing elements right if the target spot is occupied. The brute-force approach (using a dynamic array like Python's list and calling insert() each time) works for small n, but with n up to 10⁵, this hits O(n²) time complexity and times out hard.

The Key Insight: Work Backwards #

Instead of simulating each insertion step (which forces us to shift elements repeatedly), we can figure out the final position of each element directly. Here's why this works:

• When processing elements in reverse order, every element we place will not be shifted by any subsequent insertions (since there are none left).
• We just need to find the index[i]-th available spot in the final array (since earlier insertions will have taken up some positions).

How to Track Available Positions #

To efficiently find and mark available positions, we can use a segment tree that keeps track of the number of available slots in each interval. This lets us:

1. Query the k-th available position in O(log n) time.
2. Mark a position as used (unavailable) in O(log n) time.

Step-by-Step Walkthrough (Using the Example) #

Let's use your sample input to see how this works:

• arr = [1, 2, 3, 4, 5], index = [0, 1, 2, 1, 2], final array size = 5.

We process elements from last to first:

1. Element 5, index=2: We need the 3rd available slot (since index is 0-based). Available slots are [0,1,2,3,4] → pick slot 2. Mark 2 as used. Result so far: [_, _, 5, _, _]
2. Element 4, index=1: Need the 2nd available slot. Available slots are [0,1,3,4] → pick slot 1. Mark 1 as used. Result: [_, 4, 5, _, _]
3. Element 3, index=2: Need the 3rd available slot. Available slots are [0,3,4] → pick slot4. Mark4 as used. Result: [_,4,5,_,3]
4. Element2, index=1: Need the 2nd available slot. Available slots are [0,3] → pick slot3. Mark3 as used. Result: [_,4,5,2,3]
5. Element1, index=0: Need the 1st available slot. Only slot0 left. Result: [1,4,5,2,3]

Perfect—matches the sample output!

Code Implementation (Python) #

Here's a segment tree-based solution that handles 1e5 elements efficiently:

class SegmentTree:
    def __init__(self, size):
        self.n = 1
        while self.n < size:
            self.n <<= 1
        self.tree = [0] * (2 * self.n)
        # Initialize leaves: 1 means available
        for i in range(size):
            self.tree[self.n + i] = 1
        # Build the tree
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[2*i] + self.tree[2*i+1]
    
    def query_kth(self, k):
        # Find the k-th available position (1-based)
        node = 1
        while node < self.n:
            left = self.tree[2*node]
            if left >= k:
                node = 2*node
            else:
                k -= left
                node = 2*node +1
        return node - self.n
    
    def update(self, pos):
        # Mark position as unavailable (set to 0)
        pos += self.n
        self.tree[pos] = 0
        pos >>= 1
        while pos >=1:
            new_val = self.tree[2*pos] + self.tree[2*pos+1]
            if self.tree[pos] == new_val:
                break
            self.tree[pos] = new_val
            pos >>=1

def insert_elements(arr, index):
    n = len(arr)
    st = SegmentTree(n)
    result = [0]*n
    # Process in reverse order
    for i in range(n-1, -1, -1):
        # Convert 0-based index to 1-based k
        k = index[i] +1
        pos = st.query_kth(k)
        result[pos] = arr[i]
        st.update(pos)
    return result

# Test with sample input
arr = [1,2,3,4,5]
index = [0,1,2,1,2]
print(insert_elements(arr, index))  # Output: [1,4,5,2,3]

Time Complexity Analysis #

• Building the segment tree takes O(n) time.
• Each query and update operation takes O(log n) time.
• With n elements processed, total time complexity is O(n log n), which is perfect for n up to 10⁵.

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