cp-algo/number_theory/dirichlet.hpp

• View this file on GitHub
• Last update: 2026-10-16 14:16:08+02: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 <bit>
#include <iostream>
#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::has_single_bit(Align), "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;

        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);
                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>>;
}


#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 <bit>
#include <iostream>
#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::has_single_bit(Align),"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; 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); 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>>; }
#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