优化求和代码

Question

来自采访：

int fn(int a, int b)
{
    int sum = 0; 
    for (int i = a * 4; i > 0; i--)
    {
        sum += b * i * i;
    } 
    return sum;
}

如何进一步优化这段代码？我知道有一个求和公式，但我不认为记住这样的公式是面试官想要的。那么，你会如何优化它呢？

编辑：感谢 chqrlie、faivvy、asimes 和 Ap31 的建议和回答。所以我想现在将有三种优化方法：

1. 我们可以在返回时进行，而不是在每次迭代中乘以 b。
2. 用公式替换 for 循环：n * (n + 1) * (2 * n + 1) / 6 * b。 A simple derivation
3. 使用循环展开。请看asimes发的帖子。

在这三个答案中，我可能会选择 1 和 3，因为它们可以应用于具有相似结构的所有类型的代码。您应该提到有一个公式可以用作奖励，但我怀疑公式是否是面试官想要的。

还有其他建议吗？

Answer 1

公式：1*1+2*​​2+...+n*n = n(n+1)(2n+1)/6

int fn(int a, int b)
{
    a <<= 2;
    return (a*(a + 1)*((a << 1) + 1) / 6) * b;
}

这是你想要的吗？

Answer 2

除非a 为负数，否则函数fn 计算b 乘以直到4*a 的平方和。

从1 到n 的平方和可以计算为n(n+1)(2n+1)/6。

这是一个 C 翻译：

int fn(int a, int b) {
    if (a <= 0 || b == 0) {
        return 0;
    } else {
        int n = a * 4;
        return n * (n + 1) * (2 * n + 1) / 6 * b;
    }
}

正如 Ap31 所指出的，clang 足够精明，可以检测循环优化并将原始函数转换为直接计算，但它将上述代码编译为much more compact 16 assembly instructions（原始代码为 36）。

为了避免中间结果的潜在溢出，这里有一个稍微不同的公式，它不会计算更大的中间结果：

int fn(int a, int b) {
    if (a <= 0 || b == 0) {
        return 0;
    } else {
        if (a % 3 == 0)
            return (a / 3) * (4 * a + 1) * (8 * a + 1) * b * 2;
        else
            return (4 * a + 1) * (8 * a + 1) / 3 * a * b * 2;
    }
}

如果类型 long long 大于 int，则更简单的替代方法是：

int fn(int a, int b) {
    if (a <= 0 || b == 0) {
        return 0;
    } else {
        unsigned long long n = a * 4;
        return (int)(n * (n + 1) * (2 * n + 1) / 6 * b);
    }
}

Answer 3

面试官当然希望从@faivvy（和@chqrlie）的答案中得到优化，你总是可以推导出公式或者只是说你知道它存在并且你可以完全摆脱循环。

不要忘记一些常见的错误：a 可能是负数，a*a*(2*a + 1) 可能溢出。

另外需要注意的是modern compilers can do this by themselves - 你也可以向面试官提及。

Answer 4

正如@faivvy 在他的回答中指出的那样，您可以尝试完全取消 for 循环

但是，另一种方法（正确处理负数a）是执行循环展开，我将调用该函数fnUnroll。如果您不熟悉循环展开，其想法是减少迭代次数并并行求和

正如cmets中提到的，每次迭代不需要乘以b，可以在最后完成。我添加了另一个名为 fnUnrollNoMult 的函数来显示这个

#include <chrono>
#include <cstdlib>
#include <iostream>

int fn(int a, int b) {
    int sum = 0;
    for (int i = a * 4; i > 0; i--)
        sum += b * i * i;
    return sum;
}

int fnUnroll(int a, int b) {
    // Set up some number of accumulators, I picked 4
    int sum0 = 0;
    int sum1 = 0;
    int sum2 = 0;
    int sum3 = 0;

    int i = 1;
    int limit = a * 4;

    // Sum 4 values in parallel
    for ( ; i < limit; i += 4) {
        sum0 += b * i * i;
        sum1 += b * (i + 1) * (i + 1);
        sum2 += b * (i + 2) * (i + 2);
        sum3 += b * (i + 3) * (i + 3);
    }

    // Handle the remainder (if any)
    for ( ; i < limit; i++)
        sum0 += b * i + i;

    // Sum the accumulators
    return sum0 + sum1 + sum2 + sum3;
}

int fnUnrollNoMult(int a, int b) {
    int sum0 = 0;
    int sum1 = 0;
    int sum2 = 0;
    int sum3 = 0;

    // Remove b from the loops
    int i = 1;
    int limit = a * 4;
    for ( ; i < limit; i += 4) {
        sum0 += i * i;
        sum1 += (i + 1) * (i + 1);
        sum2 += (i + 2) * (i + 2);
        sum3 += (i + 3) * (i + 3);
    }
    for ( ; i < limit; i++)
        sum0 += i + i;

    // Handle b here
    return b * (sum0 + sum1 + sum2 + sum3);
}

int main(int argc, char** argv) {
    // Expects two arguments: a and b
    if (argc != 3) {
        std::cout << "Usage: " << argv[0] << " <int> <int>\n";
        return 1;
    }

    int a = atoi(argv[1]);
    int b = atoi(argv[2]);

    // This is just to demonstrate correctness
    for (int i = 0; i < 100; i++)
        for (int j = 0; j < 100; j++)
            if (
                fn(i, j) != fnUnroll(i, j) ||
                fn(i, j) != fnUnrollNoMult(i, j)
            ) {
                std::cout << "Not equal: " << i << ", " << j << std::endl;
                return 1;
            }

    // Benchmark
    using namespace std::chrono;
    {
        auto start = high_resolution_clock::now();
        int result = fn(a, b);
        auto stop  = high_resolution_clock::now();
        std::cout << "fn value:             " << result << std::endl;
        std::cout << "fn nanos:             " << duration_cast<nanoseconds>(stop - start).count() << std::endl;
    }
    {
        auto start = high_resolution_clock::now();
        int result = fnUnroll(a, b);
        auto stop  = high_resolution_clock::now();
        std::cout << "fnUnroll value:       " << result << std::endl;
        std::cout << "fnUnroll nanos:       " << duration_cast<nanoseconds>(stop - start).count() << std::endl;
    }
    {
        auto start = high_resolution_clock::now();
        int result = fnUnrollNoMult(a, b);
        auto stop  = high_resolution_clock::now();
        std::cout << "fnUnrollNoMult value: " << result << std::endl;
        std::cout << "fnUnrollNoMult nanos: " << duration_cast<nanoseconds>(stop - start).count() << std::endl;
    }

    return 0;
}

下面的程序需要两个参数，分别代表a 和b。下面我将程序编译为g++ -std=c++14 foo.cpp -O3，并得到一些a 值的这些结果：

./a.out 1 2
fn value:             60
fn nanos:             373
fnUnroll value:       60
fnUnroll nanos:       209
fnUnrollNoMult value: 60
fnUnrollNoMult nanos: 157
./a.out 1000 2
fn value:             -267004960
fn nanos:             3509
fnUnroll value:       -267004960
fnUnroll nanos:       2820
fnUnrollNoMult value: -267004960
fnUnrollNoMult nanos: 1568
./a.out 1000000 2
fn value:             -619707648
fn nanos:             3137685
fnUnroll value:       -619707648
fnUnroll nanos:       2387840
fnUnrollNoMult value: -619707648
fnUnrollNoMult nanos: 1220519