Recover an array from its insertion indices

Question

An array is built starting with an empty array, and then a sequence of insertions:

1. insert $a_1$ at index $z_1=1$
2. insert $a_2$ at index $z_2$
3. insert $a_3$ at index $z_3$
4. ...

and so on. When we insert element $a_i$ and index $z_i$ the result is that $a_i$ is now at index $z_i$, whereas everything before index $z_i$ is unchanged and everything after has its index increased by 1. (With one-based indexing) E.g., the sequence $(3,1), (5,2), (1,2)$ gives $[3]$ then $[3,5]$ then $[3,1,5]$. All of the instructions will make sense, i.e., $1 \leq z_i\leq i$.

My question is about how to calculate the final array. The naive approach would be to start with an empty array, and literally obey the instructions; but in the language I program, insertions have a worst case time complexity of $O(\# $elements shifted$)$; if, for example, all the insertion indices were $1$, this would result in $O(N^2)$ time.

Supposing we have access to all the instructions simultaneously, how can we calculate the final array in faster than $O(N^2)$ time? I would be happy with a name if this problem is well studied. I did make the following (family of) observations:

1. the element which ends up at index 1 was the last one which arrived with index 1.
2. the element which ends up at index 2 arrived with index 1 or 2. If it arrived with index 2, then none came after it with either index 1 or 2. If it arrived with index 1, then exactly one came after it with index 1, after which came none with index 2.
3. the element which ends up at index 3 arrived with index 1, 2, or 3. If it arrived with index 3, then none came after it with index 1, 2 or 3. If it arrived with index 2, then exactly one element with index 1 or 2 followed it, then none after that with index 1, 2, or 3. If it arrived with index 1, then it was followed by one with index 1, then one with index 1 or 2, and then none after that with index 1, 2, or 3.

... and so on. However I can't think of the algorithms or data structures that would make this information useful.

Update: There are two $O(n\log n)$ solutions below: the first one, which I accepted, that uses an AVL tree, and one which I learned of afterward that uses a semgent tree, and is somewhat simpler.

Answer 1

Prelimiaries

You can augment an AVL tree to support all the usual operations plus the following:

Shift$(a,b)$ increases all keys $k \ge a$ by $b \ge 0$ in $O(\log n)$ time (where $n$ is the number of elements in the tree).

To do so add a value $x_v$ to each node $v$ in the tree. This value represents an offset to be added to all keys stored in the subtree rooted at $v$. The Search, Insert, and Shift operations, along with the required rotations can be implemented as follows (I won't be using the Delete operation, but it can also be implemented).

Search The search operation work as usual except that you now keep track of the cumulative offset in the path from the current node to the root.

Insert To insert a node with key $k$, use the search operation to find the position where a node with key $k$ would need to be placed and the cumulative offset $\overline{x}$ up to that point. Add a leaf in that position and store its key as $k - \overline{x}$. Perform the necessary rotations to rebalance the tree (see the sequel).

Rotations To perform a right rotation on $u$ let $v$ be its left child. "Push down" the offset of $u$ as follows: increment the stored key of $u$ by $x_u$, add $x_u$ to the offsets of the children of $u$, and set $x_u$ to $0$. Similarly, "push down" the offset of $v$. Perform the rotation as usual. Left rotations are symmetric.

Shift$(a,b)$. Find the node $u$ with key $a$ or, if no such node exists, find its successor (if the successor doesn't exist either, we are done). Increase the stored key of $u$ by $b$. If $u$ has a right child $v$ then increase $x_v$ by $b$ as well. Walk from $u$ to the root of the tree. Every time you walk to a vertex $w$ from its left child, increase the key of $w$ by $b$ and the offset $x_z$ of the right child $z$ of $w$ by $b$ (if $z$ exists).

---

Solving your problem

Keep an augmented AVL tree $T$ and consider the operations one at a time. At the end of the generic $i$-th step, the tree will contain $i$ nodes that collectively store the elements of the first $i$ operations. Each node $u$ is associated with one element of the array. The key of $u$ is exactly the position of $u$'s element in the array, as of the $i$-th operation, while the element's value is stored as satellite data in $u$.

When the operation $(a_i, z_i)$ is to be processed do a Shift$(z_i, 1)$ operation on $T$. Then, insert a new node with key $z_i$ and satellite data $a_i$ in $T$.

At the end of the process you can traverse the tree and recover the final position (the node's key) of each array element (the node's satellite data).

The total time required is $O(n \log n)$.

Answer 2

This is an explanation of the correctness of Steven's answer. Of course I would be interested if anyone came up with something simpler but it is a great solution as it stands.

Define the effective key at a node to be the node's label + the offsets of all nodes on the (inclusive) path from itself to the root. Our inductive hypothesis is that after processing the $i$-th insertion, we have $i$ nodes with effective keys $1,2,\ldots, i$, in a balanced BST ordered with respect to effective keys. Physically, the effective key for the node associated with element $a$ equals the current index of element $a$ in our array.

Algorithm, at step $(a_i,z_i)$. By hypothesis, the tree of size $i-1$ has nodes with effective keys $1$ through $i-1$. In the case $z_i=i$ we just locate the unique node with effective key $i-1$ and insert $i$ with offset $0$ as a right child. In the other case, $z_i<i$, we perform SHIFT($z_i,1$); the way that it is defined means that SHIFT$(z_i,1)$ increases the effective key of nodes $z_i, z_i+1, \ldots, i-1$ by precisely 1 and leaves all lesser keys unchanged. To insert the effective key $z_i$, we either put it in the left slot of $z_i+1$, or the rightmost descendant of $(z_i+1).left$; in either case, we store it with label $z_i-(\text{cumulative offset})$ and offset 0. Both SHIFT and this insertion happen in $O(\log N)$ time. Now we have effective keys $1$ through $i$ reprsented in a BST.

So insertions do the right thing to effective keys. Now to keep the BST balanced (and so achieve $\log n$ complexity) AVL and red-black trees use rotations. It is known that rotations preserve the orderings in a BST. The problem here is that the effective keys are calculated using the offsets leading up to the root. A rotation may change the offsets seen on the way from a node to the root, since it may change the path from that node to the root. For instance if the node $v$ is rotated up into $u$'s position, then the offset $x_u$ that would have been counted in the path from root to $v$, is no longer encountered, while the offset $x_v$ would now be counted in the calculation for $u$. To see how we might do rotations without messing up the effective keys, we introduce the "push-down" operation on a node $u$. Let $v_1, v_2$ be its children. The push down on $u$ increments $x_{v_1}$ and $x_{v_2}$ by $x_u$; increments $u$ by $x_u$; and lastly sets $x_u=0$.

During a rotation (let's say right rotation for definiteness) about node $u$ with left child $v$, we first push down $u$, then push down $v$, and finally rotate as usual: it can be seen by checking all modified parts of the tree, that the effective keys remain the same after rotation: (in the above picture, we'd check that the effective keys corresponding to nodes $v,u$ and subtrees $R,S,T$ are unchanged).

Answer 3

I recently learned a weird trick to solve the problem in $O(n\log n)$. It appeared in a comment on a codeforces blog and seems to be well(?) known among competitive programmers. It puzzles me why it works, but I will update this answer with a proof if and when I figure it out. Below is a description of the procedure: (Update: proof below)

We are sequentially inserting $n$ elements, with the $i$ element at index $z_i$ and trying to find the final indices. Initialize an array $s[\cdot]$ of length $n$, all zeros. Begin with the $n$-th (final) element and work our way backwards to the 1st element. To start, find the $z_n$-th "$0$" of $s$, counted from the right, and change it from $0$ to $1$. The final position of the $n$-th element is just $z_n$, since no elements come after to "bump" its position. We say its "offset" is $0$.

In general, for the $i$-th element; we locate the $z_i$-th "$0$" of $s$ from the right and change it to "$1$"; let the $i$-th offset equal the $\#$ of $1$s strictly to the right of the new "$1$". Then the final index of the $i$-th element is $z_i$ plus its offset.

Example: Insertions at 1,2,1,1. Initialize $s=[1,1,1,1]$

• Insertion of the 4th element at index 1: $s=[0,0,0,0] \rightarrow [0,0,0,1]$. Offset $= 0$, since there are no $1$s to the right. Final index $=1+0$.
• Insertion of the 3rd element at index 1: $s=[0,0,0,1]\rightarrow [0,0,1,1]$. Offset $=1$, since there was one $1$ to the right. Final index $=1+1$.
• Insertion of the 2nd element at index 2: $s=[0,0,1,1]\rightarrow [1,0,1,1]$. Offset $= 2$, since there were two $1$s to the right. Final index $=2+2$.
• Insertion of the 1st element at index 1: $s=[1,0,1,1]\rightarrow [1,1,1,1]$. Offset $= 2$, since there were two $1$s to the right. Final index = 3.

Using a segment tree (or possibly a fenwick tree), queries of the type "count the number of $1$s in the range $[k,n]$ of $s$" is a standard range-sum query and can be answered in $O(\log n)$.

Queries of the type "find the $j$-th $0$ in $s$ from the right", can also be answered in $O(\log n)$ time with what is essentially a binary search of the tree starting at the root. First suppose that we have a segment tree with every element of $s$ flipped ($0$s as $1$s and vice versa). In this case we could find the $j$-th $1$ from the right by identifying the rightmost $[k,n]$ such that $s[k]+s[k+1]+\cdots+s[n]=j$. The latter problem can be accomplished via a binary search starting from the root. And in this case the desired $1$ is found at $s[k]$.

Finally, the updates to individual elements of $s$ (changing a $0$ to a $1$) can be done in $O(\log n)$ time.

---

Proof of correctness: Denote $b_i:=$ the number of elements which bump the $i$-th element; observe that every bump causes the index to rise by $1$. So the final index of the $i$-th element is $z_i+b_i$.

Lemma: Let $j>i$. Then the $j$-th element bumps the $i$ iff $z_j+b_j<z_i+b_i$.

The direction $\impliedby$ is basically obvious due to the observation about the final index; for $\implies$, assuming $j$ bumps $i$, we recognize two cases: (1) $z_j\leq z_i$; note that anything that bumps $j$ also bumps $i$; furthermore $j$ bumps $i$, so we have $b_i\geq b_j+1$, hence the result. (2) $z_j>z_i$; in this case we must have that the $i$-th element was bumped at least $z_j-z_i$ times before seeing the $j$-th element (otherwise the $j$-th element would have been inserted to the right of the $i$-th). Then the $j$-th element arrives, bumping the $i$-th. Then anything that bumps the $j$-th element also bumps the $i$-th. So we have $b_i \geq z_j-z_i+1+b_j$, proving the result.

Now initialize an array $S'$ of length $n$ with all $0$s. Starting from $i=n$ and going down to $i=1$, we set $S'[z_i+b_i] := 1$. And by the lemma above, we have that $b_i = \#\{k>z_i+b_i \colon S[k]=1\}$.

Now in the algorithm, define $d_i =$ the offset (i.e. the "number of $1$s to the right" at the step where we process the $i$-th element). Observe that $b_n = d_n$ and that otherwise, the sequences $b_i$ and $d_i$, from $i=n$ to $1$, are built following the same rules; in particular setting the $z_j$-th zero of $s$, counted from the right, is equivalent to setting the $(z_j+d_j)$-th element of $s$. This shows that $d_i=b_i$ for all $i$.