diff --git a/Data/Map.hs b/Data/Map.hs
index 399f74c7abf4ffa67cefb265e47fa5ab876d86d1..b8fcf7179bab0d732c6d5cf9fdf97ae9f9ca63b6 100644
--- a/Data/Map.hs
+++ b/Data/Map.hs
@@ -523,8 +523,8 @@ updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
 updateAt f i Tip  = error "Map.updateAt: index out of range"
 updateAt f i (Bin sx kx x l r)
   = case compare i sizeL of
-      LT -> updateAt f i l
-      GT -> updateAt f (i-sizeL-1) r
+      LT -> balance kx x (updateAt f i l) r
+      GT -> balance kx x l (updateAt f (i-sizeL-1) r)
       EQ -> case f kx x of
               Just x' -> Bin sx kx x' l r
               Nothing -> glue l r