This is an exercise in Apostol's Introduction to Analytic Number Theory about the prime number theorem (Chapter 13 Exercise 2). The problem follows:
Let $a(n)$ be an arithmetical function and $A(x)=\sum_{n\leqslant x}a(n)$ where $$a(n)=\begin{cases} 0, & n\text{ is not a prime power}\\ \frac{1}{k}, &n=p^k \end{cases}$$ Prove that $A(x)=\pi(x)+O(\sqrt{x}\log\log x)$.
I personally could approach this problem in two ways. First, we could write $$A(x)=\sum_{m=1}^\infty\frac{1}{m}\sum_{p\leqslant x^{1/m}}1=\sum_{m=1}^{[\log_2 m]}\frac{\pi(x^{1/m})}{m}$$ By the asymptotic formula of $\pi$, we obtain $$A(x)=\pi(x)+\sum_{m=2}^{[\log_2 x]}\frac{x^{1/m}+o(x^{1/m})}{\frac{m}{m}\log x}=\pi(x)+O(\sqrt{x}/\log 2)=\pi(x)+O(\sqrt{x})$$ Since $O(\sqrt{x})=O(\sqrt{x}\log\log x)$, we have completed the proof. However, I don't believe this is the correct way of proving the statement. I tried to expand $A(x)$ in another way but it seemed not promising $$\begin{align*} A(x) &=\sum_{p\leqslant x}\sum_{m\leqslant \frac{\log x}{\log p}}\frac{1}{m}\\ &=\sum_{p\leqslant x}\left(\log \left(\frac{\log x}{\log p}\right)+\gamma+O\left(\frac{\log p}{\log x}\right)\right)\\ &=\pi(x)\log \log x-\sum_{p\leqslant x}\log\log p+\gamma\pi(x)+O\left(\frac{\vartheta(x)}{\log x}\right) \end{align*}$$ This looks even worse than the first one. Is there a way that we can directly produce the error term $O(\sqrt{x}\log\log x)$? Or else there is no need to actually write $O(\sqrt{x}\log\log x)$ instead of $O(\sqrt{x})$.
There is indeed a very similar question here. But the answer to this question only provides an argument that is the same as my first argument. It did not explain why we write $O(\sqrt{x}\log\log x)$ instead of $O(\sqrt{x})$, nor did it give a direct approach to the error term.