fps/fps_sqrt.hpp

• View this file on GitHub
• Last update: 2025-06-04 11:53:52+09:00
• Include: #include "fps/fps_sqrt.hpp"

Depends on

• math_mod/mod_sqrt.hpp
• modint/mont_arb.hpp
• type_traits/integral.hpp
• type_traits/io.hpp

Verified with

• verify/yosupo_fps/fps_sprase_sqrt.test.cpp
• verify/yosupo_fps/fps_sqrt.test.cpp

Code

#ifndef KK2_FPS_FPS_SQRT_HPP
#define KK2_FPS_FPS_SQRT_HPP 1

#include <cassert>

#include "../math_mod/mod_sqrt.hpp"

namespace kk2 {

template <class FPS, class mint = typename FPS::value_type> FPS sqrt(const FPS &f, int deg = -1) {
    // using mint = typename FPS::value_type;
    if (deg == -1) deg = (int)f.size();
    if ((int)f.size() == 0) return FPS(deg, mint(0));
    if (f[0] == mint(0)) {
        for (int i = 1; i < (int)f.size(); i++) {
            if (f[i] != mint(0)) {
                if (i & 1) return {};
                if (deg - i / 2 <= 0) break;
                auto ret = sqrt(f >> i, deg - i / 2);
                if (ret.empty()) return {};
                ret = ret << (i / 2);
                if ((int)ret.size() < deg) ret.resize(deg, mint(0));
                return ret;
            }
        }
        return FPS(deg, mint(0));
    }

    long long sqr = mod_sqrt(f[0].val(), mint::getmod());
    if (sqr == -1) return {};
    assert(sqr * sqr % mint::getmod() == f[0].val());
    FPS ret = {mint(sqr)};
    mint inv2 = mint(2).inv();
    for (int i = 1; i < deg; i <<= 1) { ret = (ret + f.pre(i << 1) * ret.inv(i << 1)) * inv2; }
    return ret.pre(deg);
}

template <class FPS, class mint = typename FPS::value_type>
FPS sparse_sqrt(const FPS &f, int deg = -1) {
    // using mint = typename FPS::value_type;
    if (deg == -1) deg = (int)f.size();
    if ((int)f.size() == 0) return FPS(deg, mint(0));
    if (f[0] == mint(0)) {
        for (int i = 1; i < (int)f.size(); i++) {
            if (f[i] != mint(0)) {
                if (i & 1) return {};
                if (deg - i / 2 <= 0) break;
                auto ret = sparse_sqrt(f >> i, deg - i / 2);
                if (ret.empty()) return {};
                ret = ret << (i / 2);
                if ((int)ret.size() < deg) ret.resize(deg, mint(0));
                return ret;
            }
        }
        return FPS(deg, mint(0));
    }
    long long sqr = mod_sqrt(f[0].val(), mint::getmod());
    if (sqr == -1) return {};
    return f.sparse_pow(((mint::getmod() + 1) >> 1), deg) * mint(sqr).inv();
}

} // namespace kk2

#endif // KK2_FPS_FPS_SQRT_HPP

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/cplusplus.py", line 187, in bundle
    bundler.update(path)
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 401, in update
    self.update(self._resolve(pathlib.Path(included), included_from=path))
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 401, in update
    self.update(self._resolve(pathlib.Path(included), included_from=path))
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 401, in update
    self.update(self._resolve(pathlib.Path(included), included_from=path))
  File "/opt/hostedtoolcache/Python/3.12.0/x64/lib/python3.12/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 312, in update
    raise BundleErrorAt(path, i + 1, "#pragma once found in a non-first line")
onlinejudge_verify.languages.cplusplus_bundle.BundleErrorAt: type_traits/integral.hpp: line 4: #pragma once found in a non-first line

Back to top page