Codeforces Round #Pi (Div. 2) Problem C

Codeforces Round #Pi (Div. 2) Problem C
題目大意：給你n個數a[1],a[2],……a[n](− 109 ≤ a[i] ≤ 109) ，和一個k（1 ≤ n, k ≤ 2⋅105 ），在這n個數中找到3個數，a[i], a[j], a[x]，使這三個數成等比數列，公比是k，且三個數的下標：1<=i<j<x<=n ，三個數不必連續，問能夠找到多少這樣的組合。
既然是等比數列則有：a[i]=a[j]/k，a[j]=a[x]/k，a[i]=a[x]/k/k 。
很容易想到的一種最簡單暴力的一種做法就是直接枚舉3個數的位置，時間複雜度O(n3) ，明顯超時，需要一種高效的算法。
可以這樣想，去枚舉三個數中間的那個數，設爲a[j]，既然是中間的那個數，那麼這個a[j] % k=0，在j之前找到a[i]=a[j]/k ，記這個a[i]個數爲cnt1個，在j之後找到a[x]=a[j]∗k ，記個a[x]個數爲cnt2，那麼以a[j]爲中間的那個數，就能夠得到cnt1*cnt2中組個使得這三個數成等比數列，並且公比是k，且下標是嚴格遞增的。
有了這個想法之後，想想還是得枚舉中間的a[j],在去前面找a[i]，後面找a[x]，時間複雜度還是沒變啊，其實可以巧妙地用一種方法。
C++的STL中有兩個庫函數lower_bound()，upper_bound()。lower_bound(begin,end,val)用於在數組或者容器中，範圍是[begin,end)找到一個最小可以插入val的位置,upper_bound(begin,end,val)則是找到一個最大可以插入的位置。
所以，可以這樣處理，將每個數記錄其初始的位置，然後以其值排序，就可以通過lower_bound()和upper_bound()，來高效地枚舉了。
注意爆int啊啊啊啊。
代碼如下，僅供參考：

/*
Author:Royecode
Date:2015-08-06
*/
#include <bits/stdc++.h>
#define Pii pair <ll, int>
#define ll long long
using namespace std;
const int MAXN = 200007;
Pii arr[MAXN];
int main()
{
    ll n, k;
    while(~scanf("%I64d%I64d", &n, &k))
    {
        for(int i = 0; i < n; ++i)
        {
            scanf("%I64d", &arr[i].first);
            arr[i].second = i;
        }
        sort(arr, arr + n);
        ll ans = 0;
        for(int i = 1; i < n - 1; ++i)
        {
            if(arr[i].first % k == 0)//枚舉中間的數
            {
                //arr[i]/k的個數
                ll cnt1 = upper_bound(arr, arr + n, Pii(arr[i].first / k, arr[i].second - 1)) - lower_bound(arr, arr + n, Pii(arr[i].first / k, 0));
                //arr[i]×k的個數
                ll cnt2 = upper_bound(arr, arr + n, Pii(arr[i].first * k, n)) - lower_bound(arr, arr + n, Pii(arr[i].first * k, arr[i].second + 1));
                ans += cnt1 * cnt2;
            }
        }
        printf("%I64d\n", ans);
    }
    return 0;
}
/*
input:
5 2
1 1 2 2 4
10 3
1 2 6 2 3 6 9 18 3 9
output:
4 
6
*/

時間複雜度：O(nlogn) 。