Show that a function $f: A \to B$ is surjective if and only if it has a right inverse. 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?
 A: I think you start having difficulty right around this statement:

By the axiom of choice, I can select a single element $p\in \mathbb{R}$ and simply let $g(b) = p$.

That is not the point of the axiom of choice. The axiom of choice means, intuitively, that we can make infinitely many choices all at once.
As such, your proof doesn't really highlight how to use the AC.
I'd phrase it as follows. Recall that AC has one of its formulations as follows.
If your textbook gave a different definition, it may be a good exercise to prove it is equivalent to mine.

Axiom of Choice. Let $\mathscr{C}$ be a nonempty set, not containing the empty set. Then there exists a function $F : \mathscr{C} \to \bigcup \mathscr{C}$ such that $F(S) \in S$ for all $S \in \mathscr{C}$.

(Recall that $\bigcup \mathscr{C}$ is the union of all sets in $\mathscr{C}$, i.e., $\bigcup \mathscr{C} = \bigcup_{S \in \mathscr{C}} S$.)
In other words, AC guarantees the existence of a function that, for a collection of sets, 'picks' an element from each set in that collection.
Now suppose that a function $f : A \to B$ is surjective. Then define the collection
$$ \mathscr{C} := \{ f^{-1}(\{b\}) \subseteq A \mid b\in B \}.$$
(Recall that $f^{-1}(\{b\})$ is the set of all $a \in A$ such that $f(a) = b$.)
Since $f$ is surjective, none of the sets in $\mathscr{C}$ is empty. Thus we may apply  AC to it to obtain a map
$$ F : \mathscr{C} \to A$$
which picks one element from each set in $\mathscr{C}$ (note that $\bigcup \mathscr{C} = A$). Now define
$$ g : B \to A : b \mapsto F( f^{-1}(\{b\})).$$
That is, when given a $b \in B$, we consider the set $f^{-1}(\{b\})$ (which is nonempty) and 'pick' an element from it. This makes $g$ a right-inverse for $f$. As an exercise, try to write out why this is so.
A: The axiom of choice is equivalent to Zorn’s Lemma that is more useful to prove a lot of results.
You can consider this set
$\Lambda:=\{(h,C): C\subseteq B, h: C\to A , f\circ h=Id_C\}$
This set is non-empty because if you consider $a\in A$ then
$(h, \{f(a)\})\in \Lambda \neq \emptyset$ where $h(f(a))=a$
You can fix a particular order $<$ on $\Lambda$ :
$(h,C)< (r,S)$  if and only if $ C\subseteq S$ and $r_{|C}=h$
You can verify that $<$ a partial order on $\Lambda$.
Now we consider a chain $X\subset \Lambda$
Then $(\cup_{(h,C)\in X} h, \cup_{(h,C)\in X} C)\in \Lambda $ and it is a maggiorante for the chain $X$.
So by Zorn’s lemma there exists a maximal element $(h,C)$  for $(\Lambda,<)$.
You can prove that $C=B$ in the hypothesis that $f$ is surjective.
If there exists an element $b\in B/C$ by surjectivity of $f$ there exists an element $a\in A$ such that $f(a)=b$.
Then you can define $r: C\cup \{b\} \to A$ such that $r(x)=h(x)$ if $x\in C$ and $r(x)=a$ if $x=f(a)$.
So $(r,C\cup \{b\})\in \Lambda$ and $(h,C)<(r,C\cup \{b\})$ and it is not possible by maximality of $(h,C)$
A: The use of the axiom of choice actually comes in after the following sentence in your question: 

"...This shows that for any element $x$ of $B$, I can find a single candidate for $g(x)\in A$ such that $f(g(x))=x$...."

You can do that observation for any fixed $x\in A$ but that on its own does not assure yet the existence of a function $g:B\to A$  that sends every $y\in B$ to a suitable candidate.
For that last step the axiom of choice comes in. 
One of its forms is: 

If $\mathcal P$ is a partition of $A$ then a function $\nu:\mathcal P\to A$ exists such that $\nu(P)\in P$ for every $P\in\mathcal P$. 

This can be applied here on the partition $\mathcal P:=\{f^{-1}(\{b\})\mid b\in B\}$.
Note that this $\mathcal P$ is indeed a partition of $A$ because $f$ is surjective ensuring that $f^{-1}(\{b\})\neq\varnothing$ for every $b\in B$. 
If $h:B\to\mathcal P$ is prescribed by $b\mapsto f^{-1}(\{b\})$ then for $g$ we can take the composition $\nu\circ h:B\to A$.
Also this way of working makes a split up in two cases redundant.
