Let $A$ be a nonempty set and let $f:A\rightarrow B$ be a function. Prove that $f$ is one to one if and only if there exists a function $g:B\rightarrow A$ such that $g \circ f=i_a$ ($g \circ f$ is function composition and $i_a$ is the identity function.
Proof: ($\Rightarrow$) Assume that $f$ is one to one and $A$ is a nonempty set. Let set $C=$range $f$. Since f is one to one then every $b\in C$ there is a unique element call it $a\in A$ such that $f(a)=b$. Now we define a function $g$ where $g:B\rightarrow A$ as such: Since $C\subseteq B$ it follows that $|C|\leq |B|$. Then for each $a\in A$ we let $a=g(b)$ for some $b\in C$. Thus $g(b)=g(f(a))=a=1_A$. Since $C\subseteq B$ any element(s) call it $k$ where $k\in B$\ $C$ we can assign any $a_0\in A$ such that $g(k)=a_0$ since $k$ is not mapped by $f$.
($\Leftarrow$) Assume there exists a function $g:B\rightarrow A$ such that $g \circ f=i_a$ Let $a_1,a_2 \in A$ such that $a_1\neq a_2$ Then it follows that $g(f(a_1))=a_1$ and $g(f(a_2))=a_2$ Thus $g(f(a_1))\neq g(f(a_2))$ Thus $g \circ f$ is one to one. Let $x_1,x_2\in A$. Then it follow that $g(f(x_1))=g(f(x_2)) \implies$ that $x_1=x_2$ by definition of one to one. Thus it follow that $f(x_1)=f(x_2)$ Hence f is one to one.
How could I better this proof or what parts need fixing?