Lean4
/-- O(log n). Insert an element into the set, preserving balance and the BST property.
If an equivalent element is already in the set, the function `f` is used to generate
the element to insert (being passed the current value in the set).
insertWith f 0 {1, 2, 3} = {0, 1, 2, 3}
insertWith f 1 {1, 2, 3} = {f 1, 2, 3}
Using a preorder on `ℕ × ℕ` that only compares the first coordinate:
insertWith f (1, 1) {(0, 1), (1, 2)} = {(0, 1), f (1, 2)}
insertWith f (3, 1) {(0, 1), (1, 2)} = {(0, 1), (1, 2), (3, 1)} -/
def insertWith (f : α → α) (x : α) : Ordnode α → Ordnode α
| nil => ι x
| node sz l y r =>
match cmpLE x y with
| Ordering.lt => balanceL (insertWith f x l) y r
| Ordering.eq => node sz l (f y) r
| Ordering.gt => balanceR l y (insertWith f x r)
