You are on the right approach to attach "more data" to a balanced search tree.
For all $x\in A\cup B$, let $$d(x)=\left | \left \{ y\in A: y<x \right \} \right | - \left | \left \{ y\in B: y>x \right \} \right |,$$ which is an increasing function on $x$.
For simplicity, let us assume all elements are distinct first. We can use a balanced search tree $T$ to store all elements in $A$ and $B$, keeping track of whether each element comes from $A$ or $B$.
- We will also store $(x_\ell, \delta_\ell)$ , where $x_\ell$ is the largest element $x$ in $T$ such that $d(x)\le0$ and where $\delta_\ell=d(x_\ell)$.
- We will also store $(x_r, \delta_r)$ , where $x_r$ is the smallest element $x$ in $T$ such that $d(x)\ge0$ and where $\delta_r=d(x_r)$.
Either pair might not exist. However, at least one pair must be available, which is enough to identify the wanted $x\in A\cup B$, i.e., $d(x)=0$, if exists.
Suppose $v$ will be inserted into $T$. Here is a sketch of what to do to maintain the data structure. The idea is simply to move $x_\ell$ or $x_r$ to its neighbor one or two times. Several obvious boundary cases such as when $d(x)>0$ for all $x$ in $T$ before or after $v$ is inserted have been skipped.
- Insert $v$.
Update $(x_\ell, \delta_\ell)$ as the following.
- Do the following.
- find $x_{-1}$, the predecessor of $x_\ell$. Compute $d(x_{-1})$.
- find $x_{-2}$, the predecessor of $x_{-1}$. Compute $d(x_{-2})$.
- find $x_{1}$, the successor of $x_\ell$. Compute $d(x_1)$.
- find $x_{2}$, the successor of $x_1$. Compute $d(x_2)$.
- Let $x_{\text{new}}$ be the largest element $x$ among $x_{-2}, x_{-1}, x_\ell, x_{1}, x_{2}$ such that $d(x)\le0$. Update $(x_\ell, \delta_\ell)$ to $(x_{\text{new}}, d(x_{\text{new}}))$
Update $(x_r, \delta_r)$ similarly, which should take $O(\log n)$ time as well.
Any kind of balanced search trees can be used as long as the insertion, finding the predecessor and finding the successor can be done in $O(\log n)$. In particular, we can use AVL tree or red-black tree.
- Step 1 takes $O(\log n)$ time.
- Each search for predecessor or successor in Step 2 takes $O(\log n)$ time. Each computation of $d(*)$ takes $O(1)$ time, as we know $d(x_\ell)$. For example, $d(x_{-1})=d(x_\ell)-\lambda_A(x_{-1})-\lambda_B(x_\ell)$, where $\lambda_A(x_{-1})$ is 1 if $x_{-1}\in A$ and 0 otherwise. $\lambda_B(x_{\ell})$ is 1 if $x_{\ell}\in B$ and 0 otherwise. Similarly, $d(x_1)=d(x_\ell)+\lambda_A(x_\ell)+\lambda_B(x_1)$. So step 2 takes $O(\log n)$ tims.
- Similarly, step 3 takes $O(\log n)$ time.
When there are duplicated elements in $A$ and $B$, the idea is similar but somewhat more complicated. A balanced search tree $T$ will store all distinct elements in $A$ and $B$. Each node will store
- a distinct element $x\in A\cup B$,
- how many time that element appears in $A$,
- and how many time that element appears in $B$.
The order of the node is determined by the order of that element. Attached to the whole $T$ are the two pairs $(x_\ell, \delta_\ell)$ and $(x_r, \delta_r)$ defined the same as above.
Exercise 1. Refine the algorithm so that at any point of time, only one of the pairs $(x_\ell, \delta_\ell)$ and $(x_r, \delta_r)$ is maintained.
Exercise 2. Refine the algorithm so that only 1 or 2 cases in step 2.1 are needed by considering $v<x_\ell$ or not, $v=x_{-1}$ or not, $v\in A$ or $v\in B$, etc.
Exercise 3. Consider using one balanced search tree for elements in $A$ and another one for elements in $B$ instead of using one for both of them. Will it work?
