Codeforces Round #Pi (Div. 2) C. Geometric Progression

原题: http://codeforces.com/contest/567/problem/C

---

题目大意: 给一个n大小数组, 问其中的子序列(非连续)能组成多少个已p为公比的等比序列, 且长度为3.

---

分析: 这题卡了好久, 我用找等差数列的方法改的. 后来总是过不了test. 可是怎么看我的是O(nlgn)的复杂度, 所以过不了大的test case 肯定是有O(n)的. 于是开始想啊想啊….最后想到了是用等比中项的性质.

首先上一个过不了的code:

public void solve(int testNumber, InputReader in, OutputWriter out) {
        int n = in.readInt();
        int p = in.readInt();
        IntPair[] nums = new IntPair[n];
        for (int i = 0; i < nums.length; i++) {
            nums[i] = new IntPair(in.readInt(),i);
        }
        Arrays.sort(nums);
        int count = 0;
        for (int i = 1; i < nums.length-1; i++) {
            int j = i - 1;
            int k = i + 1;
            while (j >= 0 && k < nums.length) {
                if (nums[j].first * nums[k].first > nums[i].first*nums[i].first)
                    j--;
                else if (nums[j].first * nums[k].first < nums[i].first*nums[i].first)
                    k++;
                else {
                    if ((nums[i].first/nums[j].first) != p && (nums[k].first/nums[i].first) != p)
                        continue;
                        count++;
                        while (j > 0 && nums[j].first == nums[j - 1].first && nums[j].second < nums[i].second && nums[i].second <nums[k].second) {
                            j--;
                            count++;
                        }
                        while (k < nums.length - 1 && nums[k].first == nums[k + 1].first && nums[j].second < nums[i].second && nums[i].second <nums[k].second) {
                            k++;
                            count++;
                        }
                    j--;
                    k++;
                }
            }
        }
        out.print(count);
    }
}

最后看了一下答案, 发现是用两个set保存当前数字为等比中项的前后序列, 然后记录一下相同项的个数. 第二次扫描的时候, 如果当前数字能整除公比, 那么找到j,i,k的个数,相乘就是答案. 这里注意要用long存. 难怪我test老挂

答案code:

public void solve(int testNumber, InputReader in, OutputWriter out) {
        int n = in.readInt();
        int p = in.readInt();
        Map<Long,Long> l = new HashMap<>();
        Map<Long,Long> r = new HashMap<>();
        long res = 0;
        long[] nums = new long[n];
        for (int i = 0; i < n; i++) {
            nums[i] = in.readLong();
            if (r.containsKey(nums[i]))
                r.put(nums[i], r.get(nums[i])+1);
            else
                r.put(nums[i],1l);
            l.put(nums[i],0l);
        }
        for (int i = 0; i < n; i++) {
            r.put(nums[i],r.get(nums[i])-1l);
            if (nums[i] % p == 0){
                if (l.containsKey(nums[i]/p) && r.containsKey(nums[i]*p))
                    res += l.get(nums[i]/p) * r.get(nums[i]*p);
            }
            l.put(nums[i],l.get(nums[i])+1l);
        }
        out.print(res);
    }