C++で素数判定！

AtCoderのARC044-Aで｢素数判定｣という問題を解きました。
A: 素数判定 - AtCoder Regular Contest 044 | AtCoder

解答の中で初めて素数判定のコードを書いて教育的だったのでコードを残します。
※ARC-Aの解答コードではないです。

#include<bits/stdc++.h>
typedef long long int ll;
using namespace std;
int main(){
    ll n; cin >> n;
    bool flag = true;
    if(n==1){cout << "Not Prime" << endl; return 0;}
    for(ll i=2; i*i<=n; i++){
        if(n%i==0){
	    flag = false;
	    break;
	}
    }
    cout << (flag?"Prime":"Not Prime") << endl;
    return 0;
}

実は初め、約数をカウントして約数の数が2だったら素数...
みたいな判定をしてたのですが…
TLEしてしまい...√nまでの探索で計算量が少なく済む方法を教えて頂きました。

この素数判定を書いた後に、実際のA問題のその他の条件を付加していったのですが、
コード自体は難しい所がないのに何故かつらかった...おしまい。