cp-algo/number_theory/dirichlet.hpp

• View this file on GitHub
• Last update: 2025-12-12 22:29:00+01:00
• Include: #include "cp-algo/number_theory/dirichlet.hpp"

Depends on

• cp-algo/util/big_alloc.hpp

Verified with

• Dirichlet Inverse and Prefix Sums (verify/number_theory/dirichlet_inv.test.cpp)
• Dirichlet Convolution and Prefix Sums (verify/number_theory/dirichlet_mul.test.cpp)
• Enumerate Quotients (verify/number_theory/quotients.test.cpp)
• Sum of Totient Function (verify/number_theory/totient_sum.test.cpp)

Code

#ifndef CP_ALGO_NUMBER_THEORY_DIRICHLET_HPP
#define CP_ALGO_NUMBER_THEORY_DIRICHLET_HPP
#include "../util/big_alloc.hpp"
#include <algorithm>
#include <cstdint>
#include <ranges>
#include <vector>
#include <cmath>
namespace cp_algo::math {    
    auto floor_stats(int64_t n) {
        auto rt_n = int(sqrtl(n));
        return std::pair{rt_n, 2 * rt_n - (n / rt_n == rt_n)};
    }

    auto floor_generator(int64_t n) {
        auto [rt_n, num_floors] = floor_stats(n);
        return [n, rt_n = rt_n, num_floors = num_floors](int k) {
            return k <= rt_n ? int64_t(k) : n / int64_t(num_floors - k + 1);
        };
    }

    auto floors(int64_t n) {
        auto [_, m] = floor_stats(n);
        return std::views::iota(0, m+1) | std::views::transform(floor_generator(n));
    }

    struct interval {
        int lo, hi;
        auto operator <=>(const interval&) const = default;
    };

    // callback(k) when:
    //     (F * G)[k] = H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1]
    // Return the value to be saved in H[k]
    enum exec_mode { standard, reverse };
    template<exec_mode mode = standard>
    void Dirichlet_helper(int64_t n, auto &H, auto const& F, auto const& G, auto &&callback) {
        auto [rt_n, num_floors] = floor_stats(n);

        auto to_ord = [&](int64_t k) {
            return k <= rt_n ? int(k) : num_floors - int(n / k) + 1;
        };

        auto call = [&](interval x, interval y, interval z) {
            auto Fx = F[x.hi] - F[x.lo - 1];
            auto Fy = F[y.hi] - F[y.lo - 1];
            decltype(Fx) Gx, Gy;
            if constexpr (mode == standard) {
                Gy = G[y.hi] - G[y.lo - 1];
                Gx = G[x.hi] - G[x.lo - 1];
            } else {
                Gy = G[y.lo - 1] - G[y.hi];
                Gx = G[x.lo - 1] - G[x.hi];
            }
            auto t = Fx * Gy;
            if(x != y) [[likely]] {
                t += Fy * Gx;
            }
            H[z.lo] += t;
            if (z.hi < num_floors) [[likely]] {
                H[z.hi + 1] -= t;
            }
        };
        for (int k = 2; k <= num_floors; ++k) {
            if(k > rt_n) {
                int z = num_floors - k + 1;
                for (int x = 2; ; x++) {
                    int y_lo_ord = std::max(x, z) + 1;
                    int y_hi_ord = to_ord(n / (x * z));
                    if (y_hi_ord < y_lo_ord) break;
                    call({x, x}, {y_lo_ord, y_hi_ord}, {k, k});
                }
            }

            H[k] = callback(k);

            if(k <= rt_n) {
                int x = k;
                for (int y = 2; y < k; ++y) {
                    int z_lo_ord = to_ord(1LL * x * y);
                    int z_hi_ord = to_ord(n / x);
                    if (z_hi_ord < z_lo_ord) break;
                    call({x, x}, {y, y}, {z_lo_ord, z_hi_ord});
                }
                int z_lo_ord = to_ord(1LL * x * x);
                call({x, x}, {x, x}, {z_lo_ord, num_floors});
            }
        }
    }

    auto Dirichlet_mul(auto const& F, auto const& G, int64_t n) {
        auto m = size(F);
        std::decay_t<decltype(F)> H(m);
        H[1] = F[1] * G[1];
        Dirichlet_helper(n, H, F, G, [&](auto k) {
            return H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1];
        });
        partial_sum(begin(H), end(H), begin(H));
        return H;
    }

    void Dirichlet_div_inplace(auto &H, auto const& G, int64_t n) {
        auto Gi = G[1].inv();
        H[0] -= H[0];
        adjacent_difference(begin(H), end(H), begin(H));
        H[1] *= Gi;
        Dirichlet_helper<reverse>(n, H, H, G, [&](auto k) {
            return (Gi * (H[k] - (G[k] - G[k-1]) * H[1])) + H[k-1];
        });
    }

    auto Dirichlet_div(auto const& H, auto const& G, int64_t n) {
        auto m = std::size(G);
        using elem_t = std::decay_t<decltype(*H.begin())>;
        big_vector<elem_t> F(H.begin(), H.begin() + m);
        Dirichlet_div_inplace(F, G, n);
        return F;
    }
}
#endif // CP_ALGO_NUMBER_THEORY_DIRICHLET_HPP

#line 1 "cp-algo/number_theory/dirichlet.hpp"


#line 1 "cp-algo/util/big_alloc.hpp"



#include <set>
#include <map>
#include <deque>
#include <stack>
#include <queue>
#include <vector>
#include <string>
#include <cstddef>
#include <iostream>
#include <generator>
#include <forward_list>

// Single macro to detect POSIX platforms (Linux, Unix, macOS)
#if defined(__linux__) || defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#  define CP_ALGO_USE_MMAP 1
#  include <sys/mman.h>
#else
#  define CP_ALGO_USE_MMAP 0
#endif

namespace cp_algo {
    template <typename T, size_t Align = 32>
    class big_alloc {
        static_assert( Align >= alignof(void*), "Align must be at least pointer-size");
        static_assert(std::popcount(Align) == 1, "Align must be a power of two");
    public:
        using value_type = T;
        template <class U> struct rebind { using other = big_alloc<U, Align>; };
        constexpr bool operator==(const big_alloc&) const = default;
        constexpr bool operator!=(const big_alloc&) const = default;

        big_alloc() noexcept = default;
        template <typename U, std::size_t A>
        big_alloc(const big_alloc<U, A>&) noexcept {}

        [[nodiscard]] T* allocate(std::size_t n) {
            std::size_t padded = round_up(n * sizeof(T));
            std::size_t align = std::max<std::size_t>(alignof(T),  Align);
#if CP_ALGO_USE_MMAP
            if (padded >= MEGABYTE) {
                void* raw = mmap(nullptr, padded,
                                PROT_READ | PROT_WRITE,
                                MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
                madvise(raw, padded, MADV_HUGEPAGE);
                madvise(raw, padded, MADV_POPULATE_WRITE);
                return static_cast<T*>(raw);
            }
#endif
            return static_cast<T*>(::operator new(padded, std::align_val_t(align)));
        }

        void deallocate(T* p, std::size_t n) noexcept {
            if (!p) return;
            std::size_t padded = round_up(n * sizeof(T));
            std::size_t align  = std::max<std::size_t>(alignof(T),  Align);
    #if CP_ALGO_USE_MMAP
            if (padded >= MEGABYTE) { munmap(p, padded); return; }
    #endif
            ::operator delete(p, padded, std::align_val_t(align));
        }

    private:
        static constexpr std::size_t MEGABYTE = 1 << 20;
        static constexpr std::size_t round_up(std::size_t x) noexcept {
            return (x + Align - 1) / Align * Align;
        }
    };

    template<typename T> using big_vector = std::vector<T, big_alloc<T>>;
    template<typename T> using big_basic_string = std::basic_string<T, std::char_traits<T>, big_alloc<T>>;
    template<typename T> using big_deque = std::deque<T, big_alloc<T>>;
    template<typename T> using big_stack = std::stack<T, big_deque<T>>;
    template<typename T> using big_queue = std::queue<T, big_deque<T>>;
    template<typename T> using big_priority_queue = std::priority_queue<T, big_vector<T>>;
    template<typename T> using big_forward_list = std::forward_list<T, big_alloc<T>>;
    using big_string = big_basic_string<char>;

    template<typename Key, typename Value, typename Compare = std::less<Key>>
    using big_map = std::map<Key, Value, Compare, big_alloc<std::pair<const Key, Value>>>;
    template<typename T, typename Compare = std::less<T>>
    using big_multiset = std::multiset<T, Compare, big_alloc<T>>;
    template<typename T, typename Compare = std::less<T>>
    using big_set = std::set<T, Compare, big_alloc<T>>;
    template<typename Ref, typename V = void>

    using big_generator = std::generator<Ref, V, big_alloc<std::byte>>;
}

// Deduction guide to make elements_of with big_generator default to big_alloc
namespace std::ranges {
    template<typename Ref, typename V>
    elements_of(cp_algo::big_generator<Ref, V>&&) -> elements_of<cp_algo::big_generator<Ref, V>&&, cp_algo::big_alloc<std::byte>>;
}


#line 4 "cp-algo/number_theory/dirichlet.hpp"
#include <algorithm>
#include <cstdint>
#include <ranges>
#line 8 "cp-algo/number_theory/dirichlet.hpp"
#include <cmath>
namespace cp_algo::math {    
    auto floor_stats(int64_t n) {
        auto rt_n = int(sqrtl(n));
        return std::pair{rt_n, 2 * rt_n - (n / rt_n == rt_n)};
    }

    auto floor_generator(int64_t n) {
        auto [rt_n, num_floors] = floor_stats(n);
        return [n, rt_n = rt_n, num_floors = num_floors](int k) {
            return k <= rt_n ? int64_t(k) : n / int64_t(num_floors - k + 1);
        };
    }

    auto floors(int64_t n) {
        auto [_, m] = floor_stats(n);
        return std::views::iota(0, m+1) | std::views::transform(floor_generator(n));
    }

    struct interval {
        int lo, hi;
        auto operator <=>(const interval&) const = default;
    };

    // callback(k) when:
    //     (F * G)[k] = H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1]
    // Return the value to be saved in H[k]
    enum exec_mode { standard, reverse };
    template<exec_mode mode = standard>
    void Dirichlet_helper(int64_t n, auto &H, auto const& F, auto const& G, auto &&callback) {
        auto [rt_n, num_floors] = floor_stats(n);

        auto to_ord = [&](int64_t k) {
            return k <= rt_n ? int(k) : num_floors - int(n / k) + 1;
        };

        auto call = [&](interval x, interval y, interval z) {
            auto Fx = F[x.hi] - F[x.lo - 1];
            auto Fy = F[y.hi] - F[y.lo - 1];
            decltype(Fx) Gx, Gy;
            if constexpr (mode == standard) {
                Gy = G[y.hi] - G[y.lo - 1];
                Gx = G[x.hi] - G[x.lo - 1];
            } else {
                Gy = G[y.lo - 1] - G[y.hi];
                Gx = G[x.lo - 1] - G[x.hi];
            }
            auto t = Fx * Gy;
            if(x != y) [[likely]] {
                t += Fy * Gx;
            }
            H[z.lo] += t;
            if (z.hi < num_floors) [[likely]] {
                H[z.hi + 1] -= t;
            }
        };
        for (int k = 2; k <= num_floors; ++k) {
            if(k > rt_n) {
                int z = num_floors - k + 1;
                for (int x = 2; ; x++) {
                    int y_lo_ord = std::max(x, z) + 1;
                    int y_hi_ord = to_ord(n / (x * z));
                    if (y_hi_ord < y_lo_ord) break;
                    call({x, x}, {y_lo_ord, y_hi_ord}, {k, k});
                }
            }

            H[k] = callback(k);

            if(k <= rt_n) {
                int x = k;
                for (int y = 2; y < k; ++y) {
                    int z_lo_ord = to_ord(1LL * x * y);
                    int z_hi_ord = to_ord(n / x);
                    if (z_hi_ord < z_lo_ord) break;
                    call({x, x}, {y, y}, {z_lo_ord, z_hi_ord});
                }
                int z_lo_ord = to_ord(1LL * x * x);
                call({x, x}, {x, x}, {z_lo_ord, num_floors});
            }
        }
    }

    auto Dirichlet_mul(auto const& F, auto const& G, int64_t n) {
        auto m = size(F);
        std::decay_t<decltype(F)> H(m);
        H[1] = F[1] * G[1];
        Dirichlet_helper(n, H, F, G, [&](auto k) {
            return H[k] + (F[k] - F[k-1]) * G[1] + (G[k] - G[k-1]) * F[1];
        });
        partial_sum(begin(H), end(H), begin(H));
        return H;
    }

    void Dirichlet_div_inplace(auto &H, auto const& G, int64_t n) {
        auto Gi = G[1].inv();
        H[0] -= H[0];
        adjacent_difference(begin(H), end(H), begin(H));
        H[1] *= Gi;
        Dirichlet_helper<reverse>(n, H, H, G, [&](auto k) {
            return (Gi * (H[k] - (G[k] - G[k-1]) * H[1])) + H[k-1];
        });
    }

    auto Dirichlet_div(auto const& H, auto const& G, int64_t n) {
        auto m = std::size(G);
        using elem_t = std::decay_t<decltype(*H.begin())>;
        big_vector<elem_t> F(H.begin(), H.begin() + m);
        Dirichlet_div_inplace(F, G, n);
        return F;
    }
}

#ifndef CP_ALGO_NUMBER_THEORY_DIRICHLET_HPP
#define CP_ALGO_NUMBER_THEORY_DIRICHLET_HPP
#include "../util/big_alloc.hpp"
#include <algorithm>
#include <cstdint>
#include <ranges>
#include <vector>
#include <cmath>
namespace cp_algo::math{auto floor_stats(int64_t n){auto rt_n=int(sqrtl(n));return std::pair{rt_n,2*rt_n-(n/rt_n==rt_n)};}auto floor_generator(int64_t n){auto[rt_n,num_floors]=floor_stats(n);return[n,rt_n=rt_n,num_floors=num_floors](int k){return k<=rt_n?int64_t(k):n/int64_t(num_floors-k+1);};}auto floors(int64_t n){auto[_,m]=floor_stats(n);return std::views::iota(0,m+1)|std::views::transform(floor_generator(n));}struct interval{int lo,hi;auto operator<=>(const interval&)const=default;};enum exec_mode{standard,reverse};template<exec_mode mode=standard>void Dirichlet_helper(int64_t n,auto&H,auto const&F,auto const&G,auto&&callback){auto[rt_n,num_floors]=floor_stats(n);auto to_ord=[&](int64_t k){return k<=rt_n?int(k):num_floors-int(n/k)+1;};auto call=[&](interval x,interval y,interval z){auto Fx=F[x.hi]-F[x.lo-1];auto Fy=F[y.hi]-F[y.lo-1];decltype(Fx)Gx,Gy;if constexpr(mode==standard){Gy=G[y.hi]-G[y.lo-1];Gx=G[x.hi]-G[x.lo-1];}else{Gy=G[y.lo-1]-G[y.hi];Gx=G[x.lo-1]-G[x.hi];}auto t=Fx*Gy;if(x!=y)[[likely]]{t+=Fy*Gx;}H[z.lo]+=t;if(z.hi<num_floors)[[likely]]{H[z.hi+1]-=t;}};for(int k=2;k<=num_floors;++k){if(k>rt_n){int z=num_floors-k+1;for(int x=2;;x++){int y_lo_ord=std::max(x,z)+1;int y_hi_ord=to_ord(n/(x*z));if(y_hi_ord<y_lo_ord)break;call({x,x},{y_lo_ord,y_hi_ord},{k,k});}}H[k]=callback(k);if(k<=rt_n){int x=k;for(int y=2;y<k;++y){int z_lo_ord=to_ord(1LL*x*y);int z_hi_ord=to_ord(n/x);if(z_hi_ord<z_lo_ord)break;call({x,x},{y,y},{z_lo_ord,z_hi_ord});}int z_lo_ord=to_ord(1LL*x*x);call({x,x},{x,x},{z_lo_ord,num_floors});}}}auto Dirichlet_mul(auto const&F,auto const&G,int64_t n){auto m=size(F);std::decay_t<decltype(F)>H(m);H[1]=F[1]*G[1];Dirichlet_helper(n,H,F,G,[&](auto k){return H[k]+(F[k]-F[k-1])*G[1]+(G[k]-G[k-1])*F[1];});partial_sum(begin(H),end(H),begin(H));return H;}void Dirichlet_div_inplace(auto&H,auto const&G,int64_t n){auto Gi=G[1].inv();H[0]-=H[0];adjacent_difference(begin(H),end(H),begin(H));H[1]*=Gi;Dirichlet_helper<reverse>(n,H,H,G,[&](auto k){return(Gi*(H[k]-(G[k]-G[k-1])*H[1]))+H[k-1];});}auto Dirichlet_div(auto const&H,auto const&G,int64_t n){auto m=std::size(G);using elem_t=std::decay_t<decltype(*H.begin())>;big_vector<elem_t>F(H.begin(),H.begin()+m);Dirichlet_div_inplace(F,G,n);return F;}}
#endif

#line 1 "cp-algo/number_theory/dirichlet.hpp"
#line 1 "cp-algo/util/big_alloc.hpp"
#include <set>
#include <map>
#include <deque>
#include <stack>
#include <queue>
#include <vector>
#include <string>
#include <cstddef>
#include <iostream>
#include <generator>
#include <forward_list>
#if defined(__linux__) || defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#  define CP_ALGO_USE_MMAP 1
#  include <sys/mman.h>
#else
#  define CP_ALGO_USE_MMAP 0
#endif
namespace cp_algo{template<typename T,size_t Align=32>class big_alloc{static_assert(Align>=alignof(void*),"Align must be at least pointer-size");static_assert(std::popcount(Align)==1,"Align must be a power of two");public:using value_type=T;template<class U>struct rebind{using other=big_alloc<U,Align>;};constexpr bool operator==(const big_alloc&)const=default;constexpr bool operator!=(const big_alloc&)const=default;big_alloc()noexcept=default;template<typename U,std::size_t A>big_alloc(const big_alloc<U,A>&)noexcept{}[[nodiscard]]T*allocate(std::size_t n){std::size_t padded=round_up(n*sizeof(T));std::size_t align=std::max<std::size_t>(alignof(T),Align);
#if CP_ALGO_USE_MMAP
if(padded>=MEGABYTE){void*raw=mmap(nullptr,padded,PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANONYMOUS,-1,0);madvise(raw,padded,MADV_HUGEPAGE);madvise(raw,padded,MADV_POPULATE_WRITE);return static_cast<T*>(raw);}
#endif
return static_cast<T*>(::operator new(padded,std::align_val_t(align)));}void deallocate(T*p,std::size_t n)noexcept{if(!p)return;std::size_t padded=round_up(n*sizeof(T));std::size_t align=std::max<std::size_t>(alignof(T),Align);
#if CP_ALGO_USE_MMAP
if(padded>=MEGABYTE){munmap(p,padded);return;}
#endif
::operator delete(p,padded,std::align_val_t(align));}private:static constexpr std::size_t MEGABYTE=1<<20;static constexpr std::size_t round_up(std::size_t x)noexcept{return(x+Align-1)/Align*Align;}};template<typename T>using big_vector=std::vector<T,big_alloc<T>>;template<typename T>using big_basic_string=std::basic_string<T,std::char_traits<T>,big_alloc<T>>;template<typename T>using big_deque=std::deque<T,big_alloc<T>>;template<typename T>using big_stack=std::stack<T,big_deque<T>>;template<typename T>using big_queue=std::queue<T,big_deque<T>>;template<typename T>using big_priority_queue=std::priority_queue<T,big_vector<T>>;template<typename T>using big_forward_list=std::forward_list<T,big_alloc<T>>;using big_string=big_basic_string<char>;template<typename Key,typename Value,typename Compare=std::less<Key>>using big_map=std::map<Key,Value,Compare,big_alloc<std::pair<const Key,Value>>>;template<typename T,typename Compare=std::less<T>>using big_multiset=std::multiset<T,Compare,big_alloc<T>>;template<typename T,typename Compare=std::less<T>>using big_set=std::set<T,Compare,big_alloc<T>>;template<typename Ref,typename V=void>using big_generator=std::generator<Ref,V,big_alloc<std::byte>>;}namespace std::ranges{template<typename Ref,typename V>elements_of(cp_algo::big_generator<Ref,V>&&)->elements_of<cp_algo::big_generator<Ref,V>&&,cp_algo::big_alloc<std::byte>>;}
#line 4 "cp-algo/number_theory/dirichlet.hpp"
#include <algorithm>
#include <cstdint>
#include <ranges>
#line 8 "cp-algo/number_theory/dirichlet.hpp"
#include <cmath>
namespace cp_algo::math{auto floor_stats(int64_t n){auto rt_n=int(sqrtl(n));return std::pair{rt_n,2*rt_n-(n/rt_n==rt_n)};}auto floor_generator(int64_t n){auto[rt_n,num_floors]=floor_stats(n);return[n,rt_n=rt_n,num_floors=num_floors](int k){return k<=rt_n?int64_t(k):n/int64_t(num_floors-k+1);};}auto floors(int64_t n){auto[_,m]=floor_stats(n);return std::views::iota(0,m+1)|std::views::transform(floor_generator(n));}struct interval{int lo,hi;auto operator<=>(const interval&)const=default;};enum exec_mode{standard,reverse};template<exec_mode mode=standard>void Dirichlet_helper(int64_t n,auto&H,auto const&F,auto const&G,auto&&callback){auto[rt_n,num_floors]=floor_stats(n);auto to_ord=[&](int64_t k){return k<=rt_n?int(k):num_floors-int(n/k)+1;};auto call=[&](interval x,interval y,interval z){auto Fx=F[x.hi]-F[x.lo-1];auto Fy=F[y.hi]-F[y.lo-1];decltype(Fx)Gx,Gy;if constexpr(mode==standard){Gy=G[y.hi]-G[y.lo-1];Gx=G[x.hi]-G[x.lo-1];}else{Gy=G[y.lo-1]-G[y.hi];Gx=G[x.lo-1]-G[x.hi];}auto t=Fx*Gy;if(x!=y)[[likely]]{t+=Fy*Gx;}H[z.lo]+=t;if(z.hi<num_floors)[[likely]]{H[z.hi+1]-=t;}};for(int k=2;k<=num_floors;++k){if(k>rt_n){int z=num_floors-k+1;for(int x=2;;x++){int y_lo_ord=std::max(x,z)+1;int y_hi_ord=to_ord(n/(x*z));if(y_hi_ord<y_lo_ord)break;call({x,x},{y_lo_ord,y_hi_ord},{k,k});}}H[k]=callback(k);if(k<=rt_n){int x=k;for(int y=2;y<k;++y){int z_lo_ord=to_ord(1LL*x*y);int z_hi_ord=to_ord(n/x);if(z_hi_ord<z_lo_ord)break;call({x,x},{y,y},{z_lo_ord,z_hi_ord});}int z_lo_ord=to_ord(1LL*x*x);call({x,x},{x,x},{z_lo_ord,num_floors});}}}auto Dirichlet_mul(auto const&F,auto const&G,int64_t n){auto m=size(F);std::decay_t<decltype(F)>H(m);H[1]=F[1]*G[1];Dirichlet_helper(n,H,F,G,[&](auto k){return H[k]+(F[k]-F[k-1])*G[1]+(G[k]-G[k-1])*F[1];});partial_sum(begin(H),end(H),begin(H));return H;}void Dirichlet_div_inplace(auto&H,auto const&G,int64_t n){auto Gi=G[1].inv();H[0]-=H[0];adjacent_difference(begin(H),end(H),begin(H));H[1]*=Gi;Dirichlet_helper<reverse>(n,H,H,G,[&](auto k){return(Gi*(H[k]-(G[k]-G[k-1])*H[1]))+H[k-1];});}auto Dirichlet_div(auto const&H,auto const&G,int64_t n){auto m=std::size(G);using elem_t=std::decay_t<decltype(*H.begin())>;big_vector<elem_t>F(H.begin(),H.begin()+m);Dirichlet_div_inplace(F,G,n);return F;}}

Back to top page