G1. Magic Triples (Easy Version) - 题解

比赛与标签

比赛: Codeforces Round 867 (Div. 3) 标签: brute force, data structures, math, number theory 难度: *1700

嘿，主人！来看这个魔法问题喵~

主人你好呀~！今天我们来解决一个超级有趣的魔法问题哦！ฅ^•ﻌ•^ฅ

题目是这样的：我们有一个整数序列 a，需要找出其中有多少个 "魔法三元组" (i, j, k)。一个三元组要成为魔法三元组，必须满足三个条件呐：

1. i, j, k 都是数组的合法下标。
2. 这三个下标 i, j, k 必须互不相同。
3. 存在一个正整数 b，使得 a[i] * b = a[j] 并且 a[j] * b = a[k]。

简单来说，就是数组中三个不同位置上的数字 a[i], a[j], a[k] 恰好能构成一个公比为 b 的等比数列！我们的任务就是数出所有这样排列好的 (i, j, k) 三元组有多少个。

输入：会给我们多组测试数据。每组数据先是一个整数 n 代表序列长度，然后是 n 个整数 a_1, a_2, ..., a_n。 输出：对于每组数据，输出一个整数，也就是魔法三元组的总数。

我的思路分析，请听我慢慢道来喵！

喵哈哈，这个问题看起来有点绕，但只要我们像猫咪整理毛线球一样，把它拆开来看，就一点也不难啦！

问题的核心是找到形如 (x, x*b, x*b*b) 的数值序列，其中 x = a[i], x*b = a[j], x*b*b = a[k]。

第一步，也是最重要的一步，是统计每个数字在数组 a 中出现了多少次。我们可以用一个频率数组 counts 来记录，counts[v] 就表示数字 v 出现了多少次。因为题目说了 a_i <= 10^6，所以一个大小为 10^6 + 1 的数组就够用啦，查询起来超快的说！

接下来，我们可以根据公比 b 的值，把问题分成两种情况来讨论，这样逻辑就清晰多啦！

情况一：当公比 b = 1 时（懒洋洋的情况喵~）

如果 b = 1，那么我们的等比数列就变成了 a[i] * 1 = a[j] 和 a[j] * 1 = a[k]。这意味着 a[i] = a[j] = a[k]！

也就是说，我们正在寻找三个不同的下标 i, j, k，它们指向的数值是相同的。

假设某个数值 v 在数组中出现了 c 次。我们就要从这 c 个位置中，选出 3 个不同的位置来分别作为 i, j, k。这是一个排列问题哦！

• 为 i 选择位置有 c 种方法。
• 为 j 选择位置有 c-1 种方法（因为不能和 i 相同）。
• 为 k 选择位置有 c-2 种方法（不能和 i, j 相同）。

所以，对于一个出现了 c 次的数 v，它可以构成的魔法三元组数量就是 c * (c-1) * (c-2) 个。当然，前提是 c >= 3 啦。

我们只需要遍历所有在数组中出现过的数，计算这个值并累加到总答案里就好啦！

情况二：当公比 b > 1 时（充满活力的狩猎模式！）

如果 b > 1，那么 a[i], a[j], a[k] 的值一定是严格递增的，即 a[i] < a[j] < a[k]。这意味着它们的值肯定不相同，所以我们只需要关心数值本身，不用担心下标会重复的问题。

我们可以这样做：

1. 我们来枚举等比数列的第一项 a[i] 和公比 b。
2. 我们遍历数组中所有出现过的、不重复的数，让它作为 a[i]，我们叫它 x。
3. 对于每一个 x，我们再从 b = 2 开始向上枚举公比。
4. 计算出等比数列的第二项 y = x * b 和第三项 z = y * b。
5. 剪枝优化喵！ 这是一个非常重要的步骤！因为 a_i 的最大值是 10^6，所以如果计算出的 y 或者 z 超过了这个范围，那么更大的 b 也肯定会让它们超范围。所以，一旦 z > 10^6，我们就可以停止对当前 x 的更大 b 值的搜索了，这能节省大量时间！
6. 如果 y 和 z 都在范围内，我们就利用之前建好的 counts 数组，检查 y 和 z 是否在原数组中出现过（即 counts[y] > 0 且 counts[z] > 0）。
7. 如果都出现过，那么我们就找到了一个合法的数值组合 (x, y, z)！它们能构成的魔法三元组数量就是 counts[x] * counts[y] * counts[z]。把这个数量加到我们的总答案里。

把这两种情况加起来，就是最终的答案啦！是不是很简单呢？

代码实现，一起来敲代码吧，喵！

#include <iostream>
#include <vector>
#include <algorithm>

// 使用一个全局的 counts 数组来避免在多个测试用例之间重复分配内存，喵~
// 在每个测试用例结束后，我们会把它清理干净。
const int MAX_VAL = 1000000;
std::vector<int> counts(MAX_VAL + 1, 0);

void solve() {
    int n;
    std::cin >> n;

    // 用一个 vector 存储当前测试用例中出现过的、不重复的数值
    // 这样既方便我们遍历，也方便之后高效地清理 counts 数组。
    std::vector<int> unique_a;

    for (int i = 0; i < n; ++i) {
        int val;
        std::cin >> val;
        if (counts[val] == 0) {
            unique_a.push_back(val);
        }
        counts[val]++;
    }

    long long magic_triples = 0;

    // 情况 1: 公比 b = 1
    // 这意味着 a_i = a_j = a_k。
    // 我们需要从所有值为 x 的位置中，选出 3 个不同的下标 i, j, k。
    // 如果一个值 x 出现了 c 次，那么这种有序三元组的数量就是 c * (c-1) * (c-2)。
    for (int x : unique_a) {
        if (counts[x] >= 3) {
            long long c = counts[x];
            magic_triples += c * (c - 1) * (c - 2);
        }
    }

    // 情况 2: 公比 b > 1
    // 这意味着 a_i, a_j, a_k 形成一个公比 b > 1 的等比数列。
    // 此时数值 x = a_i, y = a_j, z = a_k 必须是不同的 (x < y < z)。
    // 我们遍历数组中每个出现过的不重复值 x 作为 a_i，
    // 然后再遍历可能的整数公比 b > 1。
    std::sort(unique_a.begin(), unique_a.end()); // 排序一下，虽然不是必须的，但好习惯嘛~

    for (int x : unique_a) {
        // 使用 long long 来计算 b 和乘积，防止溢出，这很重要哦！
        for (long long b = 2; ; ++b) {
            long long y = static_cast<long long>(x) * b;
            if (y > MAX_VAL) {
                break; // 剪枝！y 太大了，后面的肯定也大
            }
            long long z = y * b;
            if (z > MAX_VAL) {
                break; // 剪枝！z 太大了，可以换下一个 x 了
            }

            // 如果 y 和 z 也都在数组中出现过，我们就找到了一个有效的数值组合！
            if (counts[y] > 0 && counts[z] > 0) {
                // 组成三元组 (i, j, k) 的方式有 counts[x] * counts[y] * counts[z] 种。
                // 因为 x, y, z 的值不同，所以它们的下标 i, j, k 也必定不同。
                magic_triples += static_cast<long long>(counts[x]) * counts[y] * counts[z];
            }
        }
    }

    std::cout << magic_triples << "\n";

    // 为下一个测试用例清理 counts 数组。
    // 我们只需要把这个测试用例中用到的计数器重置为 0 就好啦。
    for (int x : unique_a) {
        counts[x] = 0;
    }
}

int main() {
    // 加速输入输出，让程序跑得像猫一样快！
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);
    int t;
    std::cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}

复杂度分析，跑得快不快喵？

• 时间复杂度: O(N + \sqrt{MAX_VAL} \cdot \sqrt{M}) 的说。

  • 这里的 N 是输入数组的长度，M 是数组中不重复元素的个数，MAX_VAL 是 a_i 的最大值 10^6。
  • 读取输入和填充 counts 数组是 O(N)。
  • 处理 b=1 的情况是 O(M)。
  • 核心是处理 b>1 的双重循环。外层循环遍历 M 个不重复元素 x，内层循环 b 的次数大约是 sqrt(MAX_VAL / x)。总的迭代次数大致可以估算为 sum_{i=1 to M} sqrt(MAX_VAL / a_i')，其中 a_i' 是第 i 小的不重复元素。这个和的一个宽松上界是 O(sqrt(MAX_VAL) * sqrt(M))。在给定的数据范围内，这个速度是完全可以接受的！

• 空间复杂度: O(MAX_VAL) 的说。

  • 我们使用了一个全局的 counts 数组来存储所有可能数值的频率，它的大小由 MAX_VAL 决定，所以空间复杂度是 O(MAX_VAL)。unique_a 向量的大小最多是 N，所以空间被 counts 数组主导。

知识点与总结，小猫咪的笔记时间！

这次解题我们学到了很多有用的技巧呢，主人快拿小本本记下来！

1. 分类讨论大法：把一个复杂问题根据某个关键变量（这里是公比 b）的取值范围，拆分成几个简单的小问题，是算法竞赛中的王道思路喵！
2. 频率数组/哈希表：对于需要快速查询某个元素出现次数的计数问题，频率数组（当数值范围不大时）或哈希表是我们的不二之选！O(1) 的查询速度超级棒。
3. 暴力枚举与剪枝：有时候暴力枚举是可行的，但聪明的猫娘一定要学会剪枝！通过 z > MAX_VAL 这个判断，我们砍掉了大量不必要的计算，让暴力算法变得高效。
4. 多组测试数据的优化：对于多组测试数据，像代码里那样使用一个全局数组，并且只清理用过的部分，可以避免反复申请和初始化大块内存，从而优化整体性能。

希望这篇题解能帮到你哦，主人！多练习，多思考，你也能成为像我一样的算法大师喵！加油！(๑•̀ㅂ•́)و✧