Merge pull request #96 from atcoder/patch/issue95

#95: fix convolution

Author: yosupo06

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "test/unittest/googletest"]
 	path = test/unittest/googletest
 	url = https://github.com/google/googletest
+[submodule "test/benchmark/benchmark"]
+	path = test/benchmark/benchmark
+	url = https://github.com/google/benchmark

atcoder/convolution.hpp

@@ -101,25 +101,29 @@ void butterfly_inv(std::vector<mint>& a) {
     }
 }
 
-}  // namespace internal
-
 template <class mint, internal::is_static_modint_t<mint>* = nullptr>
-std::vector<mint> convolution(std::vector<mint> a, std::vector<mint> b) {
+std::vector<mint> convolution_naive(const std::vector<mint>& a, const std::vector<mint>& b) {
     int n = int(a.size()), m = int(b.size());
-    if (!n || !m) return {};
-    if (std::min(n, m) <= 60) {
-        if (n < m) {
-            std::swap(n, m);
-            std::swap(a, b);
+    std::vector<mint> ans(n + m - 1);
+    if (n < m) {
+        for (int j = 0; j < m; j++) {
+            for (int i = 0; i < n; i++) {
+                ans[i + j] += a[i] * b[j];
+            }
         }
-        std::vector<mint> ans(n + m - 1);
+    } else {
         for (int i = 0; i < n; i++) {
             for (int j = 0; j < m; j++) {
                 ans[i + j] += a[i] * b[j];
             }
         }
-        return ans;
     }
+    return ans;
+}
+
+template <class mint, internal::is_static_modint_t<mint>* = nullptr>
+std::vector<mint> convolution_fft(std::vector<mint> a, std::vector<mint> b) {
+    int n = int(a.size()), m = int(b.size());
     int z = 1 << internal::ceil_pow2(n + m - 1);
     a.resize(z);
     internal::butterfly(a);
@@ -132,7 +136,25 @@ std::vector<mint> convolution(std::vector<mint> a, std::vector<mint> b) {
     a.resize(n + m - 1);
     mint iz = mint(z).inv();
     for (int i = 0; i < n + m - 1; i++) a[i] *= iz;
-    return a;
+    return std::move(a);
+}
+
+}  // namespace internal
+
+template <class mint, internal::is_static_modint_t<mint>* = nullptr>
+std::vector<mint> convolution(std::vector<mint>&& a, std::vector<mint>&& b) {
+    int n = int(a.size()), m = int(b.size());
+    if (!n || !m) return {};
+    if (std::min(n, m) <= 60) return convolution_naive(a, b);
+    return internal::convolution_fft(a, b);
+}
+
+template <class mint, internal::is_static_modint_t<mint>* = nullptr>
+std::vector<mint> convolution(const std::vector<mint>& a, const std::vector<mint>& b) {
+    int n = int(a.size()), m = int(b.size());
+    if (!n || !m) return {};
+    if (std::min(n, m) <= 60) return convolution_naive(a, b);
+    return internal::convolution_fft(a, b);
 }
 
 template <unsigned int mod = 998244353,

test/benchmark/CMakeLists.txt

@@ -0,0 +1,23 @@
+cmake_policy(SET CMP0048 NEW)
+project(ACLibrary)
+
+cmake_minimum_required(VERSION 3.17)
+
+set(GOOGLETEST_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../unittest/googletest")
+
+if(NOT "${CMAKE_CXX_STANDARD}")
+  set(CMAKE_CXX_STANDARD 14)
+endif()
+set(CMAKE_CXX_EXTENSIONS OFF)
+
+add_compile_options(-Wall -Wextra -Wshadow -Wconversion -Wno-sign-conversion -Werror)
+
+add_subdirectory(benchmark)
+include_directories(${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR})
+include(GoogleTest)
+
+include_directories(.)
+include_directories(../../)
+
+add_executable(Convolution convolution.cpp)
+target_link_libraries(Convolution benchmark::benchmark)

test/benchmark/benchmark

@@ -0,0 +1,35 @@
+#include "atcoder/convolution"
+#include <iostream>
+
+#include "benchmark/benchmark.h"
+
+using namespace std;
+using namespace atcoder;
+using mint = modint998244353;
+
+void CONV_same_length(benchmark::State& state) {    
+    vector<mint> a(state.range(0)), b(state.range(0));
+    for (int i = 0; i < state.range(0); i++) {
+        a[i] = i + 1234;
+        b[i] = i + 5678;
+    }
+    for (auto _ : state) {
+        benchmark::DoNotOptimize(convolution(a, b));
+    }
+}
+BENCHMARK(CONV_same_length)->RangeMultiplier(2)->Range(1, 1<<20);
+BENCHMARK(CONV_same_length)->DenseRange(1, 100, 1);
+
+void CONV_long_empty(benchmark::State& state) {
+    vector<mint> a(state.range(0)), b;
+    for (int i = 0; i < state.range(0); i++) {
+        a[i] = i + 1234;
+    }
+    for (auto _ : state) {
+        benchmark::DoNotOptimize(convolution(a, b));
+        benchmark::DoNotOptimize(convolution(b, a));
+    }
+}
+BENCHMARK(CONV_long_empty)->RangeMultiplier(2)->Range(1, 1 << 20);
+
+BENCHMARK_MAIN();