AtCoder ABC #300 (ユニークビジョンプログラミングコンテスト2023 春) : G - P-smooth number

Gまでは割と順調だったんだけどね。
https://atcoder.jp/contests/abc300/tasks/abc300_g

問題

正整数N,Pが与えられる。
素因数としてP以下の素数だけを持つ正整数のうちN以下のものは何通りか。

解法

サンプルより、最大ケースでは2*10^9より大きな解になるので、愚直に列挙すると間に合わない。
そこで平方分割する。

まず小さな素数(以下では17以下の7種)だけを素因数にもつN以下の正整数をDFSなどで列挙し、昇順に並べておく。
次に、大きな素数(19以上のもの)だけを素因数にもつN以下の正整数をDFSなどで列挙しよう。
仮に現在値がAの場合、小さな素数の組み合わせはN/A以下の範囲でのみ取りえることができる。
これは最初に昇順に並べた数列を二分探索すれば、個数を高速に求めることができる。

ll N,P;
const int prime_max = 1000000;
vector<int> prime;
int NP,divp[prime_max];

void cprime(int ma) {
	if(NP) return;
	for(int i=2;i<=ma;i++) if(divp[i]==0) {
		//M[i]=NP;
		prime.push_back(i); NP++;
		for(ll j=1LL*i*i;j>=i&&j<prime_max;j+=i) if(divp[j]==0) divp[j]=i;
	}
}

vector<ll> cand;
ll ret=0;
void dfs(ll cur,int step,int ma) {
	if(cur>N) return;
	if(step>=ma) return;
	
	if(cur*prime[step]<=N) {
		cand.push_back(cur*prime[step]);
		dfs(cur*prime[step],step,ma);
	}
	dfs(cur,step+1,ma);
}
void dfs2(ll cur,int step,int ma) {
	if(cur>N) return;
	if(step>=ma) return;
	if(cur*prime[step]<=N) {
		ret+=lower_bound(ALL(cand),N/(cur*prime[step])+1)-cand.begin();
		dfs2(cur*prime[step],step,ma);
	}
	dfs2(cur,step+1,ma);
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>N>>P;
	cprime(P);
	
	cand.push_back(1);
	dfs(1,0,min(7,NP));
	sort(ALL(cand));
	
	ret+=lower_bound(ALL(cand),N+1)-cand.begin();
	dfs2(1,7,NP);
	
	cout<<ret<<endl;
}

まとめ

これはすんなり思いつけた。