TopCoder SRM 671 Div 1 Medium

BearDarts.cpp

#include <vector>
#include <map>
#include <utility>
using namespace std;

//  Binary Indexed Tree
class BIT
{
    int n;
    vector<int> v;
public:
    BIT(int n_) {
        n = 1;
        while (n < n_)
            n <<= 1;
        v = vector<int>(n);
    }
    //  a[i] += x
    void add(int i, int x) {
        for (; i<n; i|=i+1)
            v[i] += x;
    }
    //  return a[0]+a[1]+…+a[i-1]
    int sum(int i) {
        int s = 0;
        for (i--; i>=0; i=(i&(i+1))-1)
            s += v[i];
        return s;
    }
};

class BearDarts{public:
long long count( vector <int> w )
{
    int n = (int)w.size();

    map<long long, vector<pair<int, int>>> M;
    for (int i=0; i<n; i++)
    for (int j=i+1; j<n; j++)
        M[1LL*w[i]*w[j]].push_back(make_pair(i,j));

    long long ans = 0LL;
    BIT bit(n);
    for (auto m: M)
    {
        vector<pair<int, int>> &v = m.second;
        int l = (int)v.size();

        for (int i=0; i<l;)
        {
            int j;
            for (j=i; j<l && v[i].first==v[j].first; j++)
                ans += bit.sum(v[j].second) - bit.sum(v[j].first+1);
            for (j=i; j<l && v[i].first==v[j].first; j++)
                bit.add(v[j].second, 1);
            i = j;
        }

        for (auto p: m.second)
            bit.add(p.second, -1);
    }

    return ans;
}};