Problem p03065 (C++)

Buggy submission (status: Time Limit Exceeded):

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
typedef long long ll;

bool isNotPrime[100000];
int primes[100000];
int nPrime;

inline ll calc(const vector<ll> &vs, ll x, int p) {
  int n = vs.size() - 1;
  ll res = vs[0] % p;
  REP(i, n) {
    res = res * x % p;
    res = (res + vs[i + 1]) % p;
  }
  return res;
}

const int T = 1000;

bool isOK2(const vector<ll> &vs, int p) {
  int n = vs.size() - 1;
  vector<ll> vs2(min(p, n + 1), 0);
  REP(i, n + 1) { vs2[(n - i) % (p - 1)] += vs[i]; }
  REP(i, vs2.size()) {
    if (vs[i] % p != 0) {
      return false;
    }
  }
  return true;
}

bool isOK(const vector<ll> &vs, int p) {
  int n = vs.size() - 1;
  if (isOK2(vs, p)) {
    return true;
  }
  if (p < T) {
    REP(c, p) {
      if (calc(vs, c, p) != 0) {
        return false;
      }
    }
  } else {
    REP(t, T) {
      int c = rand() % p;
      // cerr << ">>" << c << endl;
      if (calc(vs, c, p) != 0) {
        return false;
      }
    }
  }
  return true;
}

// 定数項が0でない場合
set<int> solve1(const vector<ll> &vs) {
  int n = vs.size() - 1;
  vector<int> ps;
  {
    int cur = abs(vs[n]);
    REP(i, nPrime) {
      int p = primes[i];
      if (p * p > cur) {
        break;
      }
      if (cur % p == 0) {
        ps.push_back(p);
        while (cur % p == 0) {
          cur /= p;
        }
      }
    }
    if (cur > 1) {
      ps.push_back(cur);
    }
  }
  set<int> res;
  REP(i, ps.size()) {
    if (isOK(vs, ps[i])) {
      res.insert(ps[i]);
    }
  }
  return res;
}

// 定数項が0かもしれない場合
// x*(x+1)*(2x+1) が2,3で割り切れる、みたいなのをケアが必要...
// あと、 x * (普通) みたいな...そっちはsolve1でケアすれば良いか
set<int> solve2(const vector<ll> &vs) {
  int n = vs.size() - 1;
  set<int> res1;
  {
    int zeros = 0;
    REP(i, n) {
      if (vs[n - i] == 0) {
        zeros++;
      } else {
        break;
      }
    }
    if (zeros > 0) {
      vector<ll> vs2;
      REP(i, n + 1 - zeros) { vs2.push_back(vs[i]); }
      res1 = solve1(vs2);
    }
  }
  set<int> res2;
  REP(p, n + 1) {
    if (isNotPrime[p]) {
      continue;
    }
    bool ok = true;
    if (p < T) {
      REP(c, p) {
        if (calc(vs, c, p) != 0) {
          ok = false;
          break;
        }
      }
    } else {
      REP(t, T) {
        int c = rand() % p;
        if (calc(vs, c, p) != 0) {
          ok = false;
          break;
        }
      }
    }
    if (ok) {
      res2.insert(p);
    }
  }
  set<int> res;
  for (int p : res1) {
    res.insert(p);
  }
  for (int p : res2) {
    res.insert(p);
  }
  return res;
}

int main(void) {
  isNotPrime[0] = isNotPrime[1] = true;
  REP(i, 100000) {
    if (!isNotPrime[i]) {
      primes[nPrime++] = i;
      for (int j = i + i; j < 100000; j += i) {
        isNotPrime[j] = true;
      }
    }
  }
  int n;
  scanf("%d", &n);
  vector<ll> vs;
  REP(i, n + 1) {
    ll v;
    scanf("%lld", &v);
    vs.push_back(v);
  }
  set<int> res;
  if (vs[n] == 0) {
    res = solve2(vs);
  } else {
    res = solve1(vs);
  }

  for (int p : res) {
    printf("%d\n", p);
  }

  return 0;
}


Fixed submission:

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
using namespace std;
#define REP(i, n) for (int i = 0; i < (int)(n); ++i)
typedef long long ll;

bool isNotPrime[100000];
int primes[100000];
int nPrime;

inline ll calc(const vector<ll> &vs, ll x, int p) {
  int n = vs.size() - 1;
  ll res = vs[0] % p;
  REP(i, n) {
    res = res * x % p;
    res = (res + vs[i + 1]) % p;
  }
  return res;
}

const int T = 100;

bool isOK2(const vector<ll> &vs, int p) {
  int n = vs.size() - 1;
  vector<ll> vs2(min(p, n + 1), 0);
  REP(i, n + 1) { vs2[(n - i) % (p - 1)] += vs[i]; }
  REP(i, vs2.size()) {
    if (vs[i] % p != 0) {
      return false;
    }
  }
  return true;
}

bool isOK(const vector<ll> &vs, int p) {
  int n = vs.size() - 1;
  if (isOK2(vs, p)) {
    return true;
  }
  if (p < T) {
    REP(c, p) {
      if (calc(vs, c, p) != 0) {
        return false;
      }
    }
  } else {
    REP(t, T) {
      int c = rand() % p;
      // cerr << ">>" << c << endl;
      if (calc(vs, c, p) != 0) {
        return false;
      }
    }
  }
  return true;
}

// 定数項が0でない場合
set<int> solve1(const vector<ll> &vs) {
  int n = vs.size() - 1;
  vector<int> ps;
  {
    int cur = abs(vs[n]);
    REP(i, nPrime) {
      int p = primes[i];
      if (p * p > cur) {
        break;
      }
      if (cur % p == 0) {
        ps.push_back(p);
        while (cur % p == 0) {
          cur /= p;
        }
      }
    }
    if (cur > 1) {
      ps.push_back(cur);
    }
  }
  set<int> res;
  REP(i, ps.size()) {
    if (isOK(vs, ps[i])) {
      res.insert(ps[i]);
    }
  }
  return res;
}

// 定数項が0かもしれない場合
// x*(x+1)*(2x+1) が2,3で割り切れる、みたいなのをケアが必要...
// あと、 x * (普通) みたいな...そっちはsolve1でケアすれば良いか
set<int> solve2(const vector<ll> &vs) {
  int n = vs.size() - 1;
  set<int> res1;
  {
    int zeros = 0;
    REP(i, n) {
      if (vs[n - i] == 0) {
        zeros++;
      } else {
        break;
      }
    }
    if (zeros > 0) {
      vector<ll> vs2;
      REP(i, n + 1 - zeros) { vs2.push_back(vs[i]); }
      res1 = solve1(vs2);
    }
  }
  set<int> res2;
  REP(p, n + 1) {
    if (isNotPrime[p]) {
      continue;
    }
    bool ok = true;
    if (p < T) {
      REP(c, p) {
        if (calc(vs, c, p) != 0) {
          ok = false;
          break;
        }
      }
    } else {
      REP(t, T) {
        int c = rand() % p;
        if (calc(vs, c, p) != 0) {
          ok = false;
          break;
        }
      }
    }
    if (ok) {
      res2.insert(p);
    }
  }
  set<int> res;
  for (int p : res1) {
    res.insert(p);
  }
  for (int p : res2) {
    res.insert(p);
  }
  return res;
}

int main(void) {
  isNotPrime[0] = isNotPrime[1] = true;
  REP(i, 100000) {
    if (!isNotPrime[i]) {
      primes[nPrime++] = i;
      for (int j = i + i; j < 100000; j += i) {
        isNotPrime[j] = true;
      }
    }
  }
  int n;
  scanf("%d", &n);
  vector<ll> vs;
  REP(i, n + 1) {
    ll v;
    scanf("%lld", &v);
    vs.push_back(v);
  }
  set<int> res;
  if (vs[n] == 0) {
    res = solve2(vs);
  } else {
    res = solve1(vs);
  }

  for (int p : res) {
    printf("%d\n", p);
  }

  return 0;
}