SRM 671 Div2 Medium BearDartsDiv2

問題

数値列が与えられた時に abc=dとなる(a<b<c<d)組み合わせの数を求める問題.

解答

  • 愚直にやるとO(n4)の組み合わせを検証する
  • 素数が500なのでO(n4)は厳しい
  • num[i][j]をw[i]がj..nにいくつあるのかを予め計算しておく
  • 位置ai,bi,ci(ai < bi < ci)のO(n3)の組み合わせについて,条件が成り立つdの個数はnum[w[ai]w[bi]w[ci]]となるので,これを計算する
class BearDartsDiv2
{
public:
  long long count(vector <int> w)
    {
      int limit = 10000001;
      int n = w.size();
      vector<vector<int> > num(n, vector<int>(n, 0));
      map<int, int> inv;
      int i, j, k;
      for(i = 0; i < n; i++) inv[w[i]] = i;
      for(i = 0; i < n; i++){
        for(j = 0; j < n; j++){
          int sum = 0;
          for(k = j; k < n; k++){
            if (w[i] == w[k]) sum++;
          }
          num[i][j] = sum;
        }
      }
      long long res = 0;
      for(i = 0; i < n; i++){
        for(j = i + 1; j < n; j++){
          for(k = j + 1; k < n-1; k++){
            long long score = (long long)w[i] * w[j] * w[k];
            if (score >= limit) continue;
            // cout << "score=" << score;
            if (inv.find(score) == inv.end()) continue;
            // cout << "score="<< score << " inv=" << inv[score] << "]" << endl;
            res += num[inv[score]][k+1];
          }
        }
      }
      return res;
    }
};