The textbook I'm reading from defines a right-inverse as follows:
Let $f: A \to B$ be a function. A right inverse of $f$ is a function $g: B \to A$ with the property that $f \circ g = id_B$ where $id_B$ is the identity function on $B$.
Now, I'm asked to show that $f: A \to B$ has a right inverse iff it is surjective. The author notes that I may use the axiom of choice in this proof. In other words, they say to assume that given a family of disjoint nonempty subsets of a set, there is a way to choose one element in each member of the family.
Proof Attempt: Suppose $f$ has a right inverse. So, there is a function $g: B \to A$ such that $f \circ g = id_B$. Now, choose an element $b \in B$. We observe that $g(b) \in A$ and that, by hypothesis, $f(g(b)) = id_B(b) = b$. So, we have found an element of $A$ (namely $g(b)$) whose image under $f$ is $b$. Since the choice of $b$ was arbitrary, we conclude that $f$ is surjective.
For the converse, suppose that $f$ is surjective. We need to construct a function $g: B \to A$ such that $f \circ g = id_B$ for all $x \in B$. In order to do this, fix $b \in B$. Since $f$ is surjective, there exists an element $g(b) \in A$ such that $f(g(b)) = b$. However, we note that there may be more than one element of $A$ whose image is $b$ under $f$ (i.e. there may be more than one possible candidate for $g(b)$). In order to isolate a single candidate for $g(b)$, we proceed by cases.
Case 1: The function $f$ sends all elements of $A$ to $b \in B$.
In this case, we may conclude by the surjectivity of $f$ that $b$ is the only element of $B$. Therefore if we pick an element $s \in A$, we may define a function $g: B \to A$ where $g(b) = s$. Then, we simply observe that $f(g(b)) = f(s) = b = id_B(b)$. Since $b$ is the only element of $B$, we may immediately conclude that $g$ is a right inverse of $f$.
Case 2: Only the elements in a proper subset of $A$ are mapped to $b$ by the function $f$.
In this case, we can divide the set $A$ into two disjoint subsets $R$ and $S$ where $R$ is the set of all elements in $A$ that get mapped to $b$ by $f$ and $S$ is the set of all elements of $A$ that are not mapped to $b$ by $f$. By the axiom of choice, I can select a single element $p \in R$ and simply let $g(b) = p$. This shows that for any element $x$ of $B$, I can always find a single candidate for $g(x) \in A$ such that $f(g(x)) =x$. Then, we can properly define a function $g: B \to A$ that sends any $x \in B$ to the single candidate $g(x)$ that we have chosen, and, by construction, $g$ will be a right inverse of $f$. This completes the proof.
Is my proof correct? This is the first time I've done a proof requiring the axiom of choice, so I'm wondering if I've invoked it correctly. Are there any ways that I could improve the proof or details that perhaps I'm missing?