Data Structures, Algorithms, & Applications in Java
Chapter 13, Exercise 13

Let y = heap[size]. When the element in heap[i] is removed from the max heap, the size of the max heap decreases by 1 and we must insert y back into the heap. Since the position heap[i] is vacant, we try to put y into heap[i]. We can put y into heap[i] iff y <= heap[i/2] and y >= max{heap[2*i], heap[2*i + 1]} (assuming i is not the root and has two children). Notice that at most one of these two conditions can be violated by y.

When y is greater than the element in the parent position, we use the strategy used to insert an element into a max heap and follow the path from i towards the root until we find the place to put y. When y is smaller than one or both of the children of i, we use the strategy used by removeMax and follow a path from i towards the leaves until we find the place for y. The resulting code is given below.
public class MaxHeapWithRemoveElementAt extends MaxHeap
{
   // constructors
   /** create a heap with the given initial capacity */
   public MaxHeapWithRemoveElementAt(int initialCapacity)
      {super(initialCapacity);}
   
   /** create a heap with initial capacity 10 */
   public MaxHeapWithRemoveElementAt()
      {super(10);}

   
   /** remove element in heap[i] and return it
     * @return null iff heap[i] is not part of the heap */
   public Comparable removeElementAt(int i)
   {
      if (i < 1 || i > size)
         // heap elements are heap[1:size]
         return null;
   
      // handle removal of last element as special case
      if (i == size)
         return heap[size--];

      // removal of element other than last one
      Comparable x = heap[i];           // element to remove
   
      // take last element out of heap, will try to put it at heap[i]
      Comparable y = heap[size--];      // last element
   
      // see if y is <= element at parent of i
      if (i == 1 || y.compareTo(heap[i / 2]) <= 0)
      {// no problem with parent element, move down the tree
       // looking for place for y 
         int ci = 2 * i; // child of i
         while (ci <= size)
         {
            // heap[ci] should be larger child of i
            if (ci < size &&
                heap[ci].compareTo(heap[ci + 1]) < 0)
               ci++;
      
            // can we put y in heap[i]?
            if (y.compareTo(heap[ci]) >= 0)
               break;   // yes
      
            // no
            heap[i] = heap[ci]; // move child up
            i = ci;             // move down a level
            ci *= 2;
         }
      }
      else
         // y is >= all elements in subtree with root i
         // move up the tree looking for place for y
         while (i != 1 && heap[i / 2].compareTo(y) < 0)
         {// cannot put y in heap[i]
            heap[i] = heap[i / 2]; // move element down
            i /= 2;                // move to parent
         }
   
      // put y into heap[i]
      heap[i] = y;

      return x;
   }
}



The complexity is O(log n) because the number of iterations of each while loop is O(log n) and each iteration takes O(1) time.