题目描述

给你一个下标从 1 开始、长度为 n 的整数数组 nums。

对于 nums[i]，如果 n 能够被 i 整除，也就是 n % i == 0，那么称它为一个特殊元素。

请返回所有特殊元素的平方和。（力扣 2778）

这是一道非常简单的题：枚举下标，判断是否为 n 的因子，是的话就把平方加到答案里。

class Solution {
public:
    int sumOfSquares(vector<int>& nums) {
        int n = nums.size();
        int ans = 0;

        for (int i = 1; i <= n; ++i) {
            if (n % i == 0) {
                ans += nums[i - 1] * nums[i - 1];
            }
        }
        return ans;
    }
};

这里最容易出错的地方只有一个：题目下标从 1 开始，而 C++ 数组下标从 0 开始。因此判断用 i，取值用 nums[i - 1]。

这道题其实就三件事

说白了，就是下面三件事：

  1. 哪些位置要算？
  2. 这些位置怎么算？
  3. 数组变不变，要查几次？

原题的答案也很简单：

  • 要算的位置：i 是 n 的因子；
  • 算法：把 nums[i - 1] 的平方加起来；
  • 查询：数组不变，只查一次。

后面说的“扩展”，就是把这三件事分别换一下，看看原来的代码还能不能用。

一、换个条件：不一定只看因子

1. 从“因子”改成“质因子”

比如题目改成：只有当 i 是 n 的质因子时，才把它加进答案。那怎么办？很简单，平方和不用改，只改前面的判断：

bool isPrime(int x) {
    if (x < 2) return false;
    for (int d = 2; d * d <= x; ++d) {
        if (x % d == 0) return false;
    }
    return true;
}

int answer(vector<int>& nums) {
    int n = nums.size();
    int ans = 0;
    for (int i = 1; i <= n; ++i) {
        if (n % i == 0 && isPrime(i)) {
            ans += nums[i - 1] * nums[i - 1];
        }
    }
    return ans;
}

如果 n 很大，也不用急着把 1 到 n 全部筛一遍。因为我们最后只关心 n 的因子，所以可以先找因子，再判断这些因子是不是质数：

for (int d = 1; d * d <= n; ++d) {
    if (n % d != 0) continue;

    if (isPrime(d)) {
        ans += 1LL * nums[d - 1] * nums[d - 1];
    }

    int other = n / d;
    if (other != d && isPrime(other)) {
        ans += 1LL * nums[other - 1] * nums[other - 1];
    }
}

判断一个数是不是质数时，试除到 sqrt(i) 就够了。如果有很多组数据，也可以提前筛到 sqrt(max_n)，不用直接筛到 max_n。当然，数据范围不大时，直接试除也完全可以。

2. 再换一个条件：完全平方数

比如只统计既是 n 的因子、又是完全平方数的位置：

bool isSquare(long long x) {
    long long r = static_cast<long long>(sqrtl(x));
    while ((r + 1) * (r + 1) <= x) ++r;
    while (r * r > x) --r;
    return r * r == x;
}

if (n % i == 0 && isSquare(i)) {
    ans += nums[i - 1] * nums[i - 1];
}

还是原来的套路：

枚举位置 i
    如果 condition(i, n) 成立
        计算 nums[i - 1] 的贡献

你还可以把条件改成 i 和 n 互质、i 是奇数，或者让 i + 1 和 n 满足某种关系。只要能在循环里判断，外面的代码基本不用动。

3. 顺便优化一下：不要每个位置都看

原题是从 1 枚举到 n，复杂度是 O(n)。但如果我们确定只看因子，就没必要把中间那些不是因子的位置也看一遍。

因子通常成对出现：如果 i 能整除 n，那么 n / i 也是因子。因此只需要枚举到 sqrt(n)：

long long answer(vector<int>& nums) {
    int n = nums.size();
    long long ans = 0;

    for (int d = 1; d * d <= n; ++d) {
        if (n % d != 0) continue;

        ans += 1LL * nums[d - 1] * nums[d - 1];

        int other = n / d;
        if (other != d) {
            ans += 1LL * nums[other - 1] * nums[other - 1];
        }
    }
    return ans;
}

这样复杂度就从 O(n) 降到了 O(sqrt(n))。不过力扣原题的数据范围不大，直接枚举已经够用了。这个写法主要是告诉我们：特殊位置很少时，可以直接找特殊位置。

这里还有一个小细节：d 和 n / d 不一定按原来的下标顺序访问。平方和不在乎顺序，所以没问题；但如果题目问的是“第一个满足条件的元素”，或者要求交替加减，就不能直接这样写，需要重新排序，或者老老实实按下标顺序来。

二、换一种计算方式：不一定非要平方和

原题对特殊位置做的是：

ans += nums[i - 1] * nums[i - 1];

这一行也可以换成别的操作，比如：

1. 加权平方和

假设第 i 个位置的权重是 i，要求返回：

$$ \sum_{i\mid n} i \times nums[i-1]^2 $$

代码只需要把贡献改成：

long long contribution = 1LL * i * nums[i - 1] * nums[i - 1];
ans += contribution;

2. 其他常见计算

同样的特殊位置还可以用来：

  • 求最大值：ans = max(ans, nums[i - 1])；
  • 求个数：++count；
  • 求带符号的和：根据 i 的奇偶决定加还是减；
  • 求乘积：product *= nums[i - 1]；
  • 统计频次：把 nums[i - 1] 当作桶下标。

所以，遇到“满足条件的元素做统计”时，可以先套这个模板：

for (int i = 1; i <= n; ++i) {
    if (condition(i, n)) {
        ans = combine(ans, transform(nums[i - 1], i));
    }
}

前面判断哪些位置，里面算这个位置贡献多少，最后把结果合起来。题目大多数时候只是这三步中的某一步换了。

三、从查一次变成查很多次

1. 数组不变，多次查询

原题只查一次，而且查的是整个数组。如果数组不变，但现在要问很多次，比如每次问 [1, r] 的答案，就可以先做一个前缀和：

vector<long long> prefix(n + 1, 0);
for (int i = 1; i <= n; ++i) {
    long long value = 0;
    if (n % i == 0) {
        value = 1LL * nums[i - 1] * nums[i - 1];
    }
    prefix[i] = prefix[i - 1] + value;
}

// 查询 [l, r]，这里的 l、r 都是 1-based
long long rangeSum(int l, int r) {
    return prefix[r] - prefix[l - 1];
}

前缀和只需要提前算一次，之后每次查询都是 O(1)。如果特殊位置特别少，也可以只保存这些位置，再用二分查找。

2. 特殊位置离散化

假设因子位置只有 p 个，我们把位置放进 pos，把对应的平方放进 value。查询 [l, r] 时，先在 pos 里找出落在这个区间里的位置：

left  = 第一个 >= l 的位置
right = 第一个 > r 的位置
答案 = value[left..right) 的前缀和

当 p 比 n 小很多时，这样不用给整个数组建前缀和，能省一些空间。

写成代码就是：

vector<int> pos;
vector<long long> value;

for (int i = 1; i <= n; ++i) {
    if (n % i == 0) {
        pos.push_back(i);
        value.push_back(1LL * nums[i - 1] * nums[i - 1]);
    }
}

vector<long long> prefix(value.size() + 1, 0);
for (size_t i = 0; i < value.size(); ++i) {
    prefix[i + 1] = prefix[i] + value[i];
}

long long query(int l, int r) {
    int left = lower_bound(pos.begin(), pos.end(), l) - pos.begin();
    int right = upper_bound(pos.begin(), pos.end(), r) - pos.begin();
    return prefix[right] - prefix[left];
}

pos 是从小到大排好的。lower_bound 找左边界，upper_bound 找右边界，[left, right) 里面的就是区间 [l, r] 中的特殊位置。

3. 如果数组还会被修改

  • 只改一个数，查整个数组：先记录哪些位置是特殊位置，修改时只更新这个位置。
  • 只改一个数，还要查区间：可以用树状数组，或者在线段树里维护答案。
  • 一段一起修改，一段一起查询：可以用带懒标记的线段树。
  • 数组长度也会变：位置和 n 都会跟着变，原来的特殊位置可能全部失效，这就不是简单改一个数了。

这些内容原题都用不上，但有一点要记住：到底是数值变了，还是判断条件也变了？ 如果数组长度变了，n % i == 0 也会跟着变，不能只更新被修改的那个元素。

四、动态修改：改一个数和改数组长度，不是一回事

这两种情况看起来都叫“修改”，其实差别很大。

先看简单的：只把第 k 个数改成 x。这里 k 是从 1 开始的位置，所以数组下标是 k - 1。只要数组长度 n 不变，哪些位置是特殊位置也不会变，我们只更新这一项：

if (n % k == 0) {
    ans -= 1LL * nums[k - 1] * nums[k - 1];
    nums[k - 1] = x;
    ans += 1LL * x * x;
} else {
    nums[k - 1] = x;
}

但如果插入或删除元素，n 会变化，后面的下标也可能跟着变。这时就得重新考虑：

  • 哪些位置还满足 n % i == 0？
  • 原来第 i 个元素是不是已经跑到别的位置了？
  • 查询区间的左右端点有没有变化？

所以大概可以这样分：

  1. 数组不变：一次查询、多次查询，用前缀和或离散化；
  2. 只改数值：特殊位置不变，做增量更新，必要时用树状数组或线段树；
  3. 长度也变化：位置和判断条件一起变化，需要更复杂的数据结构，这里先不展开。

五、从一维数组扩展到三维数组

一维数组里我们看的是一个下标 i。换成三维数组后，就变成看三个下标 (x, y, z)。

例如有一个大小为 a × b × c 的三维数组 box。如果定义：

x 是 a 的因子
y 是 b 的因子
z 是 c 的因子

如果 (x, y, z) 同时满足 x 是 a 的因子、y 是 b 的因子、z 是 c 的因子，那么 box[x][y][z] 就是要统计的特殊元素。所有这些特殊元素的平方和就是：

$$ \sum_{x\mid a,\ y\mid b,\ z\mid c} box[x][y][z]^2 $$

直接枚举三个坐标即可：

long long answer(vector<vector<vector<int>>>& box) {
    int a = box.size();
    int b = box[0].size();
    int c = box[0][0].size();
    long long ans = 0;

    for (int x = 1; x <= a; ++x) {
        for (int y = 1; y <= b; ++y) {
            for (int z = 1; z <= c; ++z) {
                if (a % x == 0 && b % y == 0 && c % z == 0) {
                    long long value = box[x - 1][y - 1][z - 1];
                    ans += value * value;
                }
            }
        }
    }
    return ans;
}

一维、二维、三维其实没有那么大的区别，只是从枚举一个下标变成了枚举三个下标。三维暴力要看 a * b * c 个位置，复杂度是 O(a * b * c)。如果每一维只看因子，就可以用前面的办法，把复杂度降到 O(sqrt(a) * sqrt(b) * sqrt(c))。如果还要查询一个三维区域，也可以用三维前缀和，让一次查询变成 O(1)。

最后把几种情况放在一起看

扩展方向 变化点 典型技术
条件扩展 因子改成质因子、完全平方数、互质等 试除、因子枚举、筛法、预处理
计算方式扩展 平方和改成最大值、乘积、频次等 按需替换计算和合并方式
查询扩展 一次查询改成多次查询、区间查询 前缀和、二分、离散化
动态扩展 值变化，或数组长度变化 增量更新、树状数组、线段树
维度扩展 一维下标改成二维、三维坐标 多维枚举、多维前缀和

总结

这道题用暴力就能通过，不过稍微换一下条件，就能得到很多不同的问题：

  1. 先看哪些位置需要计算；
  2. 再看这些位置要怎么算；
  3. 查询多了，就考虑前缀和或离散化；
  4. 数组会变，就看特殊位置是否也跟着变；
  5. 换成二维、三维，还是枚举、判断、计算。

所以，比起记住这道题的代码，更值得记住下面这个过程：

遍历要检查的位置
    判断位置是否满足条件
    计算这个位置的贡献
    把贡献合并到答案

以后遇到类似题目，不用一上来就找高级算法。先看看题目到底改了哪一步，很多看起来不一样的题，其实只是这个过程换了一部分。