#1010005 (C++23) No.2892 Lime and Karin

提出ソース
結果

問題	No.2892 Lime and Karin
ユーザー	noya2
提出日時	2024-09-13 21:55:27
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	4,289 ms / 8,000 ms
コード長	35,615 bytes
コンパイル時間	4,602 ms
コンパイル使用メモリ	280,328 KB
実行使用メモリ	27,188 KB
最終ジャッジ日時	2024-09-13 21:56:16
合計ジャッジ時間	17,360 ms
ジャッジサーバーID （参考情報）	judge1 / judge5
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 5,248 KB
testcase_01	AC	2 ms 5,376 KB
testcase_02	AC	2 ms 5,376 KB
testcase_03	AC	2 ms 5,376 KB
testcase_04	AC	2 ms 5,376 KB
testcase_05	AC	2 ms 5,376 KB
testcase_06	AC	2 ms 5,376 KB
testcase_07	AC	2 ms 5,376 KB
testcase_08	AC	2 ms 5,376 KB
testcase_09	AC	2 ms 5,376 KB
testcase_10	AC	2 ms 5,376 KB
testcase_11	AC	2 ms 5,376 KB
testcase_12	AC	3 ms 5,376 KB
testcase_13	AC	2 ms 5,376 KB
testcase_14	AC	78 ms 5,376 KB
testcase_15	AC	75 ms 5,376 KB
testcase_16	AC	123 ms 5,376 KB
testcase_17	AC	30 ms 5,376 KB
testcase_18	AC	123 ms 5,376 KB
testcase_19	AC	171 ms 5,888 KB
testcase_20	AC	12 ms 5,376 KB
testcase_21	AC	92 ms 5,376 KB
testcase_22	AC	116 ms 5,376 KB
testcase_23	AC	49 ms 5,376 KB
testcase_24	AC	175 ms 5,888 KB
testcase_25	AC	203 ms 6,016 KB
testcase_26	AC	172 ms 6,016 KB
testcase_27	AC	174 ms 6,016 KB
testcase_28	AC	174 ms 6,144 KB
testcase_29	AC	207 ms 6,016 KB
testcase_30	AC	179 ms 6,016 KB
testcase_31	AC	173 ms 6,016 KB
testcase_32	AC	177 ms 6,016 KB
testcase_33	AC	172 ms 6,016 KB
testcase_34	AC	40 ms 6,016 KB
testcase_35	AC	43 ms 6,016 KB
testcase_36	AC	44 ms 6,016 KB
testcase_37	AC	45 ms 6,144 KB
testcase_38	AC	45 ms 5,888 KB
testcase_39	AC	759 ms 19,076 KB
testcase_40	AC	760 ms 27,188 KB
testcase_41	AC	233 ms 16,376 KB
testcase_42	AC	255 ms 23,796 KB
testcase_43	AC	244 ms 18,676 KB
testcase_44	AC	68 ms 5,376 KB
testcase_45	AC	167 ms 5,504 KB
testcase_46	AC	80 ms 5,376 KB
testcase_47	AC	104 ms 5,376 KB
testcase_48	AC	45 ms 5,376 KB
testcase_49	AC	123 ms 5,376 KB
testcase_50	AC	4,289 ms 6,016 KB
testcase_51	AC	109 ms 5,888 KB
testcase_52	AC	524 ms 6,016 KB
testcase_53	AC	250 ms 6,016 KB
testcase_54	AC	423 ms 6,016 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#line 2 "/Users/noya2/Desktop/Noya2_library/template/template.hpp"
using namespace std;

#include<bits/stdc++.h>
#line 1 "/Users/noya2/Desktop/Noya2_library/template/inout_old.hpp"
namespace noya2 {

template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p){
    os << p.first << " " << p.second;
    return os;
}
template <typename T, typename U>
istream &operator>>(istream &is, pair<T, U> &p){
    is >> p.first >> p.second;
    return is;
}

template <typename T>
ostream &operator<<(ostream &os, const vector<T> &v){
    int s = (int)v.size();
    for (int i = 0; i < s; i++) os << (i ? " " : "") << v[i];
    return os;
}
template <typename T>
istream &operator>>(istream &is, vector<T> &v){
    for (auto &x : v) is >> x;
    return is;
}

void in() {}
template <typename T, class... U>
void in(T &t, U &...u){
    cin >> t;
    in(u...);
}

void out() { cout << "\n"; }
template <typename T, class... U, char sep = ' '>
void out(const T &t, const U &...u){
    cout << t;
    if (sizeof...(u)) cout << sep;
    out(u...);
}

template<typename T>
void out(const vector<vector<T>> &vv){
    int s = (int)vv.size();
    for (int i = 0; i < s; i++) out(vv[i]);
}

struct IoSetup {
    IoSetup(){
        cin.tie(nullptr);
        ios::sync_with_stdio(false);
        cout << fixed << setprecision(15);
        cerr << fixed << setprecision(7);
    }
} iosetup_noya2;

} // namespace noya2
#line 1 "/Users/noya2/Desktop/Noya2_library/template/const.hpp"
namespace noya2{

const int iinf = 1'000'000'007;
const long long linf = 2'000'000'000'000'000'000LL;
const long long mod998 =  998244353;
const long long mod107 = 1000000007;
const long double pi = 3.14159265358979323;
const vector<int> dx = {0,1,0,-1,1,1,-1,-1};
const vector<int> dy = {1,0,-1,0,1,-1,-1,1};
const string ALP = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
const string alp = "abcdefghijklmnopqrstuvwxyz";
const string NUM = "0123456789";

void yes(){ cout << "Yes\n"; }
void no(){ cout << "No\n"; }
void YES(){ cout << "YES\n"; }
void NO(){ cout << "NO\n"; }
void yn(bool t){ t ? yes() : no(); }
void YN(bool t){ t ? YES() : NO(); }

} // namespace noya2
#line 2 "/Users/noya2/Desktop/Noya2_library/template/utils.hpp"

#line 6 "/Users/noya2/Desktop/Noya2_library/template/utils.hpp"

namespace noya2{

unsigned long long inner_binary_gcd(unsigned long long a, unsigned long long b){
    if (a == 0 || b == 0) return a + b;
    int n = __builtin_ctzll(a); a >>= n;
    int m = __builtin_ctzll(b); b >>= m;
    while (a != b) {
        int mm = __builtin_ctzll(a - b);
        bool f = a > b;
        unsigned long long c = f ? a : b;
        b = f ? b : a;
        a = (c - b) >> mm;
    }
    return a << std::min(n, m);
}

template<typename T> T gcd_fast(T a, T b){ return static_cast<T>(inner_binary_gcd(std::abs(a),std::abs(b))); }

long long sqrt_fast(long long n) {
    if (n <= 0) return 0;
    long long x = sqrt(n);
    while ((x + 1) * (x + 1) <= n) x++;
    while (x * x > n) x--;
    return x;
}

template<typename T> T floor_div(const T n, const T d) {
    assert(d != 0);
    return n / d - static_cast<T>((n ^ d) < 0 && n % d != 0);
}

template<typename T> T ceil_div(const T n, const T d) {
    assert(d != 0);
    return n / d + static_cast<T>((n ^ d) >= 0 && n % d != 0);
}

template<typename T> void uniq(std::vector<T> &v){
    std::sort(v.begin(),v.end());
    v.erase(unique(v.begin(),v.end()),v.end());
}

template <typename T, typename U> inline bool chmin(T &x, U y) { return (y < x) ? (x = y, true) : false; }

template <typename T, typename U> inline bool chmax(T &x, U y) { return (x < y) ? (x = y, true) : false; }

template<typename T> inline bool range(T l, T x, T r){ return l <= x && x < r; }

} // namespace noya2
#line 8 "/Users/noya2/Desktop/Noya2_library/template/template.hpp"

#define rep(i,n) for (int i = 0; i < (int)(n); i++)
#define repp(i,m,n) for (int i = (m); i < (int)(n); i++)
#define reb(i,n) for (int i = (int)(n-1); i >= 0; i--)
#define all(v) (v).begin(),(v).end()

using ll = long long;
using ld = long double;
using uint = unsigned int;
using ull = unsigned long long;
using pii = pair<int,int>;
using pll = pair<ll,ll>;
using pil = pair<int,ll>;
using pli = pair<ll,int>;

namespace noya2{

/*　~ (. _________ . /)　*/

}

using namespace noya2;


#line 2 "c.cpp"


#line 2 "/Users/noya2/Desktop/Noya2_library/tree/simple_tree.hpp"

#line 2 "/Users/noya2/Desktop/Noya2_library/data_structure/csr.hpp"

#line 4 "/Users/noya2/Desktop/Noya2_library/data_structure/csr.hpp"
#include<ranges>
#line 7 "/Users/noya2/Desktop/Noya2_library/data_structure/csr.hpp"

namespace noya2::internal {

template<class E>
struct csr {
    csr () {}
    csr (int _n) : n(_n) {}
    csr (int _n, int m) : n(_n){
        start.reserve(m);
        elist.reserve(m);
    }
    // ACL style constructor (do not have to call build)
    csr (int _n, const std::vector<std::pair<int,E>> &idx_elem) : n(_n), start(_n + 2), elist(idx_elem.size()) {
        for (auto &[i, e] : idx_elem){
            start[i + 2]++;
        }
        for (int i = 1; i < n; i++){
            start[i + 2] += start[i + 1];
        }
        for (auto &[i, e] : idx_elem){
            elist[start[i + 1]++] = e;
        }
        prepared = true;
    }
    int add(int idx, E elem){
        int eid = start.size();
        start.emplace_back(idx);
        elist.emplace_back(elem);
        return eid;
    }
    void build(){
        if (prepared) return ;
        int m = start.size();
        std::vector<E> nelist(m);
        std::vector<int> nstart(n + 2, 0);
        for (int i = 0; i < m; i++){
            nstart[start[i] + 2]++;
        }
        for (int i = 1; i < n; i++){
            nstart[i + 2] += nstart[i + 1];
        }
        for (int i = 0; i < m; i++){
            nelist[nstart[start[i] + 1]++] = elist[i];
        }
        swap(elist,nelist);
        swap(start,nstart);
        prepared = true;
    }
    const auto operator[](int idx) const {
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    auto operator[](int idx){
        return std::ranges::subrange(elist.begin()+start[idx],elist.begin()+start[idx+1]);
    }
    const auto operator()(int idx, int l, int r) const {
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    auto operator()(int idx, int l, int r){
        return std::ranges::subrange(elist.begin()+start[idx]+l,elist.begin()+start[idx]+r);
    }
    int n;
    std::vector<int> start;
    std::vector<E> elist;
    bool prepared = false;
};

} // namespace noya2::internal
#line 5 "/Users/noya2/Desktop/Noya2_library/tree/simple_tree.hpp"

namespace noya2 {

struct simple_tree {
    internal::csr<int> g;
    simple_tree () {}
    simple_tree (int _n) : g(_n, (_n - 1)*2) {
        if (_n == 1){
            g.build();
        }
    }
    void add_edge(int u, int v){
        g.add(u, v);
        int id = g.add(v, u);
        if (id + 1 == (g.n - 1)*2) g.build();
    }
    void input(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int u, v; cin >> u >> v;
            u -= indexed, v -= indexed;
            add_edge(u, v);
        }
    }
    void input_parents(int indexed = 1){
        for (int i = 0; i < g.n - 1; i++){
            int v; cin >> v;
            v -= indexed;
            add_edge(i + 1, v);
        }
    }
    const auto operator[](int v) const {
        return g[v];
    }
    auto operator[](int v){
        return g[v];
    }
    int size() const {
        return g.n;
    }
};

} // namespace noya2
#line 2 "/Users/noya2/Desktop/Noya2_library/tree/centroid_decomposition.hpp"

#line 4 "/Users/noya2/Desktop/Noya2_library/tree/centroid_decomposition.hpp"

namespace noya2 {

std::vector<int> centroid_decomposition(const auto &g){
    int n = g.size();
    if (n == 0){
        return {};
    }
    std::vector<int> sub(n), order;
    order.reserve(n);
    auto subtree = [&](auto sfs, int v, int f) -> void {
        sub[v] = 1;
        for (int u : g[v]){
            if (u == f) continue;
            sfs(sfs, u, v);
            sub[v] += sub[u];
        }
    };
    subtree(subtree,0,-1);
    auto fixed_root = [&](auto self, int root, int par, int cur_size) -> void {
        auto dfs = [&](auto sfs, int v, int f, int sz) -> int {
            int heavy = 0, child = -1;
            for (int u : g[v]){
                if (u == f) continue;
                if (heavy < sub[u]){
                    heavy = sub[u];
                    child = u;
                }
            }
            if (heavy > sz/2){
                int ret = sfs(sfs, child, v, sz);
                sub[v] -= ret;
                return ret;
            }
            else {
                order.emplace_back(v);
                for (int u : g[v]){
                    if (u == f) continue;
                    self(self, u, v, sub[u]);
                }
                int ret = sub[v];
                sub[v] = 0;
                return ret;
            }
        };
        while (cur_size > 0){
            cur_size -= dfs(dfs, root, par, cur_size);
        }
    };
    fixed_root(fixed_root, 0, -1, n);
    return order;
}

std::vector<int> centroid_decomposition_tree(const auto &g){
    int n = g.size();
    if (n == 0){
        return {};
    }
    std::vector<int> sub(n), par_tree(n);
    auto subtree = [&](auto sfs, int v, int f) -> void {
        sub[v] = 1;
        for (int u : g[v]){
            if (u == f) continue;
            sfs(sfs, u, v);
            sub[v] += sub[u];
        }
    };
    subtree(subtree,0,-1);
    auto fixed_root = [&](auto self, int root, int par, int cur_size, int cpre) -> void {
        auto dfs = [&](auto sfs, int v, int f, int sz) -> int {
            int heavy = 0, child = -1;
            for (int u : g[v]){
                if (u == f) continue;
                if (heavy < sub[u]){
                    heavy = sub[u];
                    child = u;
                }
            }
            if (heavy > sz/2){
                int ret = sfs(sfs, child, v, sz);
                sub[v] -= ret;
                return ret;
            }
            else {
                par_tree[v] = cpre;
                for (int u : g[v]){
                    if (u == f) continue;
                    self(self, u, v, sub[u], v);
                }
                int ret = sub[v];
                cpre = v;
                sub[v] = 0;
                return ret;
            }
        };
        while (cur_size > 0){
            cur_size -= dfs(dfs, root, par, cur_size);
        }
    };
    fixed_root(fixed_root, 0, -1, n, -1);
    return par_tree;
}

} // namespace noya2
#line 2 "/Users/noya2/Desktop/Noya2_library/utility/modint4724.hpp"

#line 2 "/Users/noya2/Desktop/Noya2_library/utility/modint.hpp"

#line 4 "/Users/noya2/Desktop/Noya2_library/utility/modint.hpp"

#line 2 "/Users/noya2/Desktop/Noya2_library/math/prime.hpp"

#line 4 "/Users/noya2/Desktop/Noya2_library/math/prime.hpp"
namespace noya2 {

constexpr long long safe_mod(long long x, long long m) {
    x %= m;
    if (x < 0) x += m;
    return x;
}

constexpr long long pow_mod_constexpr(long long x, long long n, int m) {
    if (m == 1) return 0;
    unsigned int _m = (unsigned int)(m);
    unsigned long long r = 1;
    unsigned long long y = safe_mod(x, m);
    while (n) {
        if (n & 1) r = (r * y) % _m;
        y = (y * y) % _m;
        n >>= 1;
    }
    return r;
}

constexpr bool is_prime_constexpr(int n) {
    if (n <= 1) return false;
    if (n == 2 || n == 7 || n == 61) return true;
    if (n % 2 == 0) return false;
    long long d = n - 1;
    while (d % 2 == 0) d /= 2;
    constexpr long long bases[3] = {2, 7, 61};
    for (long long a : bases) {
        long long t = d;
        long long y = pow_mod_constexpr(a, t, n);
        while (t != n - 1 && y != 1 && y != n - 1) {
            y = y * y % n;
            t <<= 1;
        }
        if (y != n - 1 && t % 2 == 0) {
            return false;
        }
    }
    return true;
}
template <int n> constexpr bool is_prime_flag = is_prime_constexpr(n);

constexpr std::pair<long long, long long> inv_gcd(long long a, long long b) {
    a = safe_mod(a, b);
    if (a == 0) return {b, 0};
    long long s = b, t = a;
    long long m0 = 0, m1 = 1;
    while (t) {
        long long u = s / t;
        s -= t * u;
        m0 -= m1 * u; 
        auto tmp = s;
        s = t;
        t = tmp;
        tmp = m0;
        m0 = m1;
        m1 = tmp;
    }
    if (m0 < 0) m0 += b / s;
    return {s, m0};
}

constexpr int primitive_root_constexpr(int m) {
    if (m == 2) return 1;
    if (m == 167772161) return 3;
    if (m == 469762049) return 3;
    if (m == 754974721) return 11;
    if (m == 998244353) return 3;
    int divs[20] = {};
    divs[0] = 2;
    int cnt = 1;
    int x = (m - 1) / 2;
    while (x % 2 == 0) x /= 2;
    for (int i = 3; (long long)(i)*i <= x; i += 2) {
        if (x % i == 0) {
            divs[cnt++] = i;
            while (x % i == 0) {
                x /= i;
            }
        }
    }
    if (x > 1) {
        divs[cnt++] = x;
    }
    for (int g = 2;; g++) {
        bool ok = true;
        for (int i = 0; i < cnt; i++) {
            if (pow_mod_constexpr(g, (m - 1) / divs[i], m) == 1) {
                ok = false;
                break;
            }
        }
        if (ok) return g;
    }
}
template <int m> constexpr int primitive_root_flag = primitive_root_constexpr(m);

// constexpr long long primitive_root_constexpr(long long m){
//     if (m == (1LL << 47) - (1LL << 24) + 1) return 3;
//     return primitive_root_constexpr(static_cast<int>(m));
// }

} // namespace noya2
#line 6 "/Users/noya2/Desktop/Noya2_library/utility/modint.hpp"

namespace noya2{

struct barrett {
    unsigned int _m;
    unsigned long long im;
    explicit barrett(unsigned int m) : _m(m), im((unsigned long long)(-1) / m + 1) {}
    unsigned int umod() const { return _m; }
    unsigned int mul(unsigned int a, unsigned int b) const {
        unsigned long long z = a;
        z *= b;
        unsigned long long x = (unsigned long long)((__uint128_t(z) * im) >> 64);
        unsigned int v = (unsigned int)(z - x * _m);
        if (_m <= v) v += _m;
        return v;
    }
};

template <int m>
struct static_modint {
    using mint = static_modint;
  public:
    static constexpr int mod() { return m; }
    static mint raw(int v) {
        mint x;
        x._v = v;
        return x;
    }
    constexpr static_modint() : _v(0) {}
    template<std::signed_integral T>
    constexpr static_modint(T v){
        long long x = (long long)(v % (long long)(umod()));
        if (x < 0) x += umod();
        _v = (unsigned int)(x);
    }
    template<std::unsigned_integral T>
    constexpr static_modint(T v){
        _v = (unsigned int)(v % umod());
    }
    constexpr unsigned int val() const { return _v; }
    mint& operator++() {
        _v++;
        if (_v == umod()) _v = 0;
        return *this;
    }
    mint& operator--() {
        if (_v == 0) _v = umod();
        _v--;
        return *this;
    }
    mint operator++(int) {
        mint result = *this;
        ++*this;
        return result;
    }
    mint operator--(int) {
        mint result = *this;
        --*this;
        return result;
    }
    constexpr mint& operator+=(const mint& rhs) {
        _v += rhs._v;
        if (_v >= umod()) _v -= umod();
        return *this;
    }
    constexpr mint& operator-=(const mint& rhs) {
        _v -= rhs._v;
        if (_v >= umod()) _v += umod();
        return *this;
    }
    constexpr mint& operator*=(const mint& rhs) {
        unsigned long long z = _v;
        z *= rhs._v;
        _v = (uint)(z % umod());
        return *this;
    }
    constexpr mint& operator/=(const mint& rhs) { return *this = *this * rhs.inv(); }
    constexpr mint operator+() const { return *this; }
    constexpr mint operator-() const { return mint() - *this; }
    constexpr mint pow(long long n) const {
        assert(0 <= n);
        mint x = *this, r = 1;
        while (n) {
            if (n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }
    constexpr mint inv() const {
        if (prime) {
            assert(_v);
            return pow(umod() - 2);
        } else {
            auto eg = inv_gcd(_v, m);
            assert(eg.first == 1);
            return eg.second;
        }
    }
    friend constexpr mint operator+(const mint& lhs, const mint& rhs) {
        return mint(lhs) += rhs;
    }
    friend constexpr mint operator-(const mint& lhs, const mint& rhs) {
        return mint(lhs) -= rhs;
    }
    friend constexpr mint operator*(const mint& lhs, const mint& rhs) {
        return mint(lhs) *= rhs;
    }
    friend constexpr mint operator/(const mint& lhs, const mint& rhs) {
        return mint(lhs) /= rhs;
    }
    friend constexpr bool operator==(const mint& lhs, const mint& rhs) {
        return lhs._v == rhs._v;
    }
    friend constexpr bool operator!=(const mint& lhs, const mint& rhs) {
        return lhs._v != rhs._v;
    }
    friend std::ostream &operator<<(std::ostream &os, const mint& p) {
        return os << p.val();
    }
    friend std::istream &operator>>(std::istream &is, mint &a) {
        long long t; is >> t;
        a = mint(t);
        return (is);
    }

  private:
    unsigned int _v;
    static constexpr unsigned int umod() { return m; }
    static constexpr bool prime = is_prime_flag<m>;
};


template <int id> struct dynamic_modint {
    using mint = dynamic_modint;
  public:
    static int mod() { return (int)(bt.umod()); }
    static void set_mod(int m) {
        assert(1 <= m);
        bt = barrett(m);
    }
    static mint raw(int v) {
        mint x;
        x._v = v;
        return x;
    }

    dynamic_modint() : _v(0) {}
    template<std::signed_integral T>
    dynamic_modint(T v){
        long long x = (long long)(v % (long long)(umod()));
        if (x < 0) x += umod();
        _v = (unsigned int)(x);
    }
    template<std::unsigned_integral T>
    dynamic_modint(T v){
        _v = (unsigned int)(v % umod());
    }
    uint val() const { return _v; }
    mint& operator++() {
        _v++;
        if (_v == umod()) _v = 0;
        return *this;
    }
    mint& operator--() {
        if (_v == 0) _v = umod();
        _v--;
        return *this;
    }
    mint operator++(int) {
        mint result = *this;
        ++*this;
        return result;
    }
    mint operator--(int) {
        mint result = *this;
        --*this;
        return result;
    }
    mint& operator+=(const mint& rhs) {
        _v += rhs._v;
        if (_v >= umod()) _v -= umod();
        return *this;
    }
    mint& operator-=(const mint& rhs) {
        _v += mod() - rhs._v;
        if (_v >= umod()) _v -= umod();
        return *this;
    }
    mint& operator*=(const mint& rhs) {
        _v = bt.mul(_v, rhs._v);
        return *this;
    }
    mint& operator/=(const mint& rhs) { return *this = *this * rhs.inv(); }
    mint operator+() const { return *this; }
    mint operator-() const { return mint() - *this; }
    mint pow(long long n) const {
        assert(0 <= n);
        mint x = *this, r = 1;
        while (n) {
            if (n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }
    mint inv() const {
        auto eg = noya2::inv_gcd(_v, mod());
        assert(eg.first == 1);
        return eg.second;
    }
    friend mint operator+(const mint& lhs, const mint& rhs) {
        return mint(lhs) += rhs;
    }
    friend mint operator-(const mint& lhs, const mint& rhs) {
        return mint(lhs) -= rhs;
    }
    friend mint operator*(const mint& lhs, const mint& rhs) {
        return mint(lhs) *= rhs;
    }
    friend mint operator/(const mint& lhs, const mint& rhs) {
        return mint(lhs) /= rhs;
    }
    friend bool operator==(const mint& lhs, const mint& rhs) {
        return lhs._v == rhs._v;
    }
    friend bool operator!=(const mint& lhs, const mint& rhs) {
        return lhs._v != rhs._v;
    }
    friend std::ostream &operator<<(std::ostream &os, const mint& p) {
        return os << p.val();
    }
    friend std::istream &operator>>(std::istream &is, mint &a) {
        long long t; is >> t;
        a = mint(t);
        return (is);
    }

  private:
    unsigned int _v;
    static barrett bt;
    static unsigned int umod() { return bt.umod(); }
};
template <int id> noya2::barrett dynamic_modint<id>::bt(998244353);

using modint998244353 = static_modint<998244353>;
using modint1000000007 = static_modint<1000000007>;
using modint = dynamic_modint<-1>;

template<typename T>
concept Modint = requires (T &a){
    T::mod();
    a.inv();
    a.val();
    a.pow(declval<int>());
};

} // namespace noya2
#line 4 "/Users/noya2/Desktop/Noya2_library/utility/modint4724.hpp"

namespace noya2 {

template<>
struct static_modint<-4724> {
    static constexpr unsigned long long mod(){
        return m;
    }
    static constexpr unsigned long long cal_mod(unsigned long long x){
        unsigned long long xu = x >> 47;
        unsigned long long xd = x & MASK47;
        unsigned long long res = (xu << 24) + xd - xu;
        if (res >= m) res -= m;
        return res;
    }
    constexpr static_modint() : _v(0) {}
    constexpr static_modint(long long x){
        if (x < 0){
            _v = cal_mod(-x);
            if (_v != 0){
                _v = m - _v;
            }
        }
        else {
            _v = cal_mod(x);
        }
    }
    constexpr static_modint(unsigned long long x){
        _v = cal_mod(x);
    }
    template<std::signed_integral T>
    constexpr static_modint(T x) : static_modint((long long)x) {}
    template<std::unsigned_integral T>
    constexpr static_modint(T x) : static_modint((unsigned long long)x) {}
    
    using modint4724 = static_modint;
    constexpr modint4724 &operator+=(const modint4724 &p){
        _v += p._v;
        if (_v >= m) _v -= m;
        return *this;
    }
    constexpr modint4724 &operator-=(const modint4724 &p){
        _v += m - p._v;
        if (_v >= m) _v -= m;
        return *this;
    }
    constexpr modint4724 &operator*=(const modint4724 &p){
        unsigned long long a = _v, b = p._v;
        unsigned long long au = a >> 24, ad = a & MASK24;
        unsigned long long bu = b >> 24, bd = b & MASK24;
        unsigned long long X = (au + ad) * (bu + bd);
        unsigned long long Y = ad * bd;
        unsigned long long Z = au * bu;
        unsigned long long A = X - Y + Z, B = Y + m4 - 2*Z;
        unsigned long long Au = A >> 23, Ad = A & MASK23;
        _v = cal_mod(((Au + Ad) << 24) + B + m - Au);
        return *this;
    }
    constexpr modint4724 &operator/=(const modint4724 &p){
        *this *= p.inv();
        return *this;
    }
    friend constexpr modint4724 operator+(const modint4724 &lhs, const modint4724 &rhs){
        return modint4724(lhs) += rhs;
    }
    friend constexpr modint4724 operator-(const modint4724 &lhs, const modint4724 &rhs){
        return modint4724(lhs) -= rhs;
    }
    friend constexpr modint4724 operator*(const modint4724 &lhs, const modint4724 &rhs){
        return modint4724(lhs) *= rhs;
    }
    friend constexpr modint4724 operator/(const modint4724 &lhs, const modint4724 &rhs){
        return modint4724(lhs) /= rhs;
    }
    constexpr modint4724 operator+() const {
        return *this;
    }
    constexpr modint4724 operator-() const {
        return modint4724() - *this;
    }
    constexpr modint4724 inv() const {
        long long a = _v, b = m, u = 1, v = 0;
        while (b > 0){
            long long t = a / b;
            std::swap(a -= t * b, b);
            std::swap(u -= t * v, v);
        }
        return modint4724(u);
    }
    constexpr modint4724 pow(long long n) const {
        modint4724 ret(1ULL), mul(_v);
        while (n != 0){
            if (n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }
    friend std::istream &operator>>(std::istream &is, modint4724 &p){
        unsigned long long x;
        is >> x;
        p = modint4724(x);
        return is;
    }
    friend std::ostream &operator<<(std::ostream &os, const modint4724 &p){
        return os << p._v;
    }
    constexpr unsigned long long val() const {
        return _v;
    }
    constexpr auto operator<=>(const modint4724 &) const = default;

  private:
    unsigned long long _v;
    static constexpr unsigned long long m = (1ULL << 47) - (1ULL << 24) + 1;
    static constexpr unsigned long long m4 = m << 2;
    static constexpr unsigned long long MASK23 = (1ULL << 23) - 1;
    static constexpr unsigned long long MASK24 = (1ULL << 24) - 1;
    static constexpr unsigned long long MASK47 = (1ULL << 47) - 1;
};
using modint4724 = static_modint<-4724>;

} // namespace noya2
#line 2 "/Users/noya2/Desktop/Noya2_library/fps/formal_power_series.hpp"

#line 4 "/Users/noya2/Desktop/Noya2_library/fps/formal_power_series.hpp"

namespace noya2{

template<typename T>
concept Field = requires (T a, T b){
    a + b; a - b; a / b; a * b;
    T(0); T(1);
};

template<class Info>
concept Fps_Info = requires {
    typename Info::value_type;
    requires Field<typename Info::value_type>;
    {Info::multiply(declval<vector<typename Info::value_type>>(),declval<vector<typename Info::value_type>>())} -> convertible_to<vector<typename Info::value_type>>;
    {Info::inv(declval<vector<typename Info::value_type>>(),declval<int>())} -> convertible_to<vector<typename Info::value_type>>;
    {Info::integral(declval<vector<typename Info::value_type>>())} -> convertible_to<vector<typename Info::value_type>>;
};

template<Fps_Info Info>
struct FormalPowerSeries : vector<typename Info::value_type> {
    using T = typename Info::value_type;
    using vector<T>::vector;
    using vector<T>::operator=;
    using FPS = FormalPowerSeries;
    FormalPowerSeries (const vector<T> &init_ = {}){ (*this) = init_; }
    void shrink(){ while (!(*this).empty() && (*this).back() == T(0)) (*this).pop_back(); }
    FPS &operator+=(const T &r){
        if ((*this).empty()) (*this).resize(1);
        (*this)[0] += r;
        return *this;
    }
    FPS &operator-=(const T &r){
        if ((*this).empty()) (*this).resize(1);
        (*this)[0] -= r;
        return *this;
    }
    FPS &operator*=(const T &r){
        for (auto &x : *this) x *= r;
        return *this;
    }
    FPS &operator/=(const T &r){
        (*this) *= T(1)/r;
        return *this;
    }
    FPS &operator<<=(const int &d){
        (*this).insert((*this).begin(),d,T(0));
        return *this;
    }
    FPS &operator>>=(const int &d){
        if ((int)(*this).size() <= d) (*this).clear();
        else (*this).erase((*this).begin(),(*this).begin()+d);
        return *this;
    }
    FPS &operator+=(const FPS &r){
        if ((*this).size() < r.size()) (*this).resize(r.size());
        for (int i = 0; i < (int)(r.size()); i++) (*this)[i] += r[i];
        return *this;
    }
    FPS &operator-=(const FPS &r){
        if ((*this).size() < r.size()) (*this).resize(r.size());
        for (int i = 0; i < (int)(r.size()); i++) (*this)[i] -= r[i];
        return *this;
    }
    FPS &operator*=(const FPS &r){
        if ((*this).empty() || r.empty()){
            (*this).clear();
            return *this;
        }
        (*this) = Info::multiply(*this,r);
        return *this;
    }
    FPS operator+(const T &r) const { return FPS(*this) += r; }
    FPS operator-(const T &r) const { return FPS(*this) -= r; }
    FPS operator*(const T &r) const { return FPS(*this) *= r; }
    FPS operator/(const T &r) const { return FPS(*this) /= r; }
    FPS operator<<(const int &d) const { return FPS(*this) <<= d; }
    FPS operator>>(const int &d) const { return FPS(*this) >>= d; }
    FPS operator+(const FPS &r) const { return FPS(*this) += r; }
    FPS operator-(const FPS &r) const { return FPS(*this) -= r; }
    FPS operator*(const FPS &r) const { return FPS(*this) *= r; }
    FPS operator+() const { return *this; }
    FPS operator-() const {
        FPS res(*this);
        for (auto &x : res) x = -x;
        return res;
    }
    T eval(const T &x) const {
        T res = T(0), w = T(1);
        for (auto &e : *this) res += e * w, w *= x;
        return res;
    }
    static FPS dot(const FPS &lhs, const FPS &rhs){
        FPS res(min(lhs.size(),rhs.size()));
        for (int i = 0; i < (int)res.size(); i++) res[i] = lhs[i] * rhs[i];
        return res;
    }
    FPS pre(int siz) const {
        FPS ret((*this).begin(), (*this).begin() + min((int)this->size(), siz));
        if ((int)ret.size() < siz) ret.resize(siz);
        return ret;
    }
    FPS rev() const {
        FPS ret(*this);
        reverse(ret.begin(), ret.end());
        return ret;
    }
    FPS diff() const {
        const int n = (int)this->size();
        FPS ret(max(0, n - 1));
        T one(1), coeff(1);
        for (int i = 1; i < n; i++) {
            ret[i - 1] = (*this)[i] * coeff;
            coeff += one;
        }
        return ret;
    }
    FPS integral() const {
        FPS ret = Info::integral(*this);
        return ret;
    }
    FPS inv(int d = -1) const {
        FPS ret = Info::inv(*this,d);
        return ret;
    }
    FPS exp(int d = -1) const {
        const int n = (*this).size();
        if (d == -1) d = n;
        FPS f = {T(1)+(*this)[0],(*this)[1]}, res = {1,(n > 1 ? (*this)[1] : 0)};
        for (int sz = 2; sz < d; sz <<= 1){
            f.insert(f.end(),(*this).begin()+min(n,sz),(*this).begin()+min(n,sz*2));
            if ((int)f.size() < sz*2) f.resize(sz*2);
            res = res * (f - res.log(2*sz));
            res.resize(sz*2);
        }
        res.resize(d);
        return res;
    }
    FPS log(int d = -1) const {
        assert(!(*this).empty() && (*this)[0] == T(1));
        if (d == -1) d = (*this).size();
        return (this->diff() * this->inv(d)).pre(d - 1).integral();
    }
};

} // namespace noya2
#line 8 "c.cpp"

namespace noya2{

consteval unsigned long long primitive_root_4724(unsigned long long p){
    if (p == modint4724::mod()){
        return 3;
    }
    throw ;
}

template<Modint mint>
struct number_theoretic_transform {
    static constexpr mint pr = primitive_root_4724(mint::mod());
    static constexpr int level = std::countr_zero(mint::mod() - 1);
    static constexpr std::array<mint,level+1> wgen(mint r){
        std::array<mint,level+1> ret;
        ret[level] = r;
        for (int i = level-1; i >= 0; i--){
            ret[i] = ret[i+1] * ret[i+1];
        }
        return ret;
    }
    static constexpr std::array<mint,level+1> wp = wgen(pr.pow((mint::mod()-1) >> level));
    static constexpr std::array<mint,level+1> wm = wgen(pr.pow((mint::mod()-1) >> level).inv());
    void fft2(std::vector<mint> &a){
        if (a.empty()) return ;
        int n = a.size();
        int k = std::countr_zero((unsigned int)(n));
        assert(n == (1 << k));
        for (int t = 1, v = 1<<(k-1), wi = k; v > 0; t <<= 1, v >>= 1, wi -= 1){
            mint ww = 1;
            int pl = 1<<wi;
            for (int j = 0; j < v; j++, ww *= wm[wi]){
                int j0 = j, j1 = j0+v;
                for (int i = 0; i < t; i++, j0 += pl, j1 += pl){
                    mint a1 = a[j1];
                    a[j1] = (a[j0] - a1) * ww;
                    a[j0] += a1;
                }
            }
        }
    }
    void ifft2(std::vector<mint> &a){
        if (a.empty()) return ;
        int n = a.size();
        int k = std::countr_zero((unsigned int)(n));
        assert(n == (1 << k));
        for (int v = 1, t = 1<<(k-1), wi = 1; t > 0; v <<= 1, t >>= 1, wi += 1){
            mint ww = 1;
            int pl = 1<<wi;
            for (int j = 0; j < v; j++, ww *= wp[wi]){
                int j0 = j, j1 = j0+v;
                for (int i = 0; i < t; i++, j0 += pl, j1 += pl){
                    mint a1 = a[j1] * ww;
                    a[j1] = a[j0] - a1;
                    a[j0] += a1;
                }
            }
        }
    }
    std::vector<mint> multiply(const std::vector<mint> &a, const std::vector<mint> &b) {
        int l = a.size() + b.size() - 1;
        if (min<int>(a.size(), b.size()) <= 40){
            std::vector<mint> s(l);
            for (int i = 0; i < (int)a.size(); i++) for (int j = 0; j < (int)b.size(); j++) s[i + j] += a[i] * b[j];
            return s;
        }
        int k = 2, M = 4;
        while (M < l) M <<= 1, ++k;
        std::vector<mint> s(M);
        for (int i = 0; i < (int)a.size(); ++i) s[i] = a[i];
        fft2(s);
        if (a.size() == b.size() && a == b) {
            for (int i = 0; i < M; ++i) s[i] *= s[i];
        }
        else {
            std::vector<mint> t(M);
            for (int i = 0; i < (int)b.size(); ++i) t[i] = b[i];
            fft2(t);
            for (int i = 0; i < M; ++i) s[i] *= t[i];
        }
        ifft2(s);
        s.resize(l);
        mint invm = mint(M).inv();
        for (int i = 0; i < l; ++i) s[i] *= invm;
        return s;
    }
};

} // namespace noya2

struct fps4724info {
    using value_type = modint4724;
    using mint = modint4724;
    static std::vector<mint> multiply(const std::vector<mint> &a, const std::vector<mint> &b){
        static number_theoretic_transform<mint> ntt;
        return ntt.multiply(a,b);
    }
    static std::vector<mint> inv(const std::vector<mint> &a, int d = -1){
        assert(false);
    }
    static std::vector<mint> integral(const std::vector<mint> &a){
        assert(false);
    }
};

using mint = modint4724;
using fps = FormalPowerSeries<fps4724info>;

void solve(){
    int n; in(n);
    simple_tree g(n);
    g.input();
    string s; in(s);
    vector<bool> done(n,false);
    mint ans = 0;
    mint i2 = mint(2).inv();
    for (int ctr : centroid_decomposition(g)){
        fps f;
        auto depth = [&](auto sfs, int v, int ff) -> int {
            int ret = (s[v] == '1' ? 1 : -1);
            int mi = 0;
            for (int u : g[v]){
                if (u == ff) continue;
                if (done[u]) continue;
                int dp = sfs(sfs,u,v);
                chmin(mi,dp);
            }
            return mi + ret;
        };
        auto dfs = [&](auto sfs, int v, int ff, int d) -> void {
            d += (s[v] == '1' ? 1 : -1);
            for (int u : g[v]){
                if (u == ff) continue;
                if (done[u]) continue;
                sfs(sfs,u,v,d);
            }
            if ((int)f.size() <= d){
                f.resize(d+1);
            }
            f[d] += 1;
        };
        fps sum, sq;
        int geta = 0;
        for (int v : g[ctr]){
            if (done[v]) continue;
            int d = depth(depth,v,ctr);
            int gg = (d < 0 ? -d : 0);
            chmax(geta,gg);
        }
        for (int v : g[ctr]){
            if (done[v]) continue;
            dfs(dfs,v,ctr,geta);
            sum += f;
            sq += f*f;
            f = {};
        }
        sq = (sum*sum - sq) * i2;
        for (int i = (s[ctr] == '1' ? 0 : 2); ; i++){
            int r = i+geta*2;
            if (r >= (int)(sq.size())) break;
            ans += sq[r];
        }
        for (int i = (s[ctr] == '1' ? 0 : 2); ; i++){
            int r = i+geta;
            if (r >= (int)(sum.size())) break;
            ans += sum[r];
        }
        if (s[ctr] == '1'){
            ans += 1;
        }
        // out(ctr,ans,geta);
        done[ctr] = true;
    }
    out(ans);
}

int main(){
    int t = 1; //in(t);
    while (t--) { solve(); }
}
yukicoder

結果

テストケース

ソースコード
