Haidora and strange arrays (Editorial)

Solution Given query values up to S_{\max}=\max s, note that for any x>1 we have \sigma(x)\ge 1+x, hence if \sigma(x)\le S_{\max} then x\le S_{\max}-1, so it is sufficient to enumerate x\in[1,S_{\max}]; compute all $\sigma (x)$ for 1\le x\le S_{\max} by the classic divisor-sum sieve (for each divisor $d$, add $d$ to all multiples $kd$), then characterize primality using \sigma(p)=p+1 (and for composite x>1 we have \sigma(x)>x+1), so the allowed values are exactly those with x=1 or \sigma(x)\ne x+1; define

for each 1\le v\le S_{\max}, and count arrays by total divisor-sum using a single DP on the sum: let f[0]=1 and for t\ge 1 let

because choosing the last element with divisor-sum $v$ (in c_v ways) leaves an independent prefix contributing t-v; then the answer to a query $s$ is $f[s]$.

Complexity \mathcal{O}(s^2).

#include <bits/stdc++.h>
using namespace std;
#define ll long long
const int M=1000000007;
ll divsum[20005];
ll countt[20005];
ll ans[10005];
int main(){
    countt[0]=1;
 for(int i=1;i<=20000;i++)
 {
     for(int j=1;j*i<=20000;j++)
     {
         divsum[i*j]+=i;
     }
     if(divsum[i]!=i+1&&divsum[i]<=10000)
     {
         countt[divsum[i]]++;
     }
 }
 ans[0]=1;
 for(int i=1;i<=10000;i++)
 {
     for(int j=1;j<=i;j++)
     {
         ans[i]+=(ans[i-j]*countt[j])%M;
         ans[i]%=M;
     }
 }
 int t;
 cin>>t;
 while(t--)
 {
     int s;
     cin>>s;
     cout<<ans[s]<<'\n';
 }
 return 0;   
}