library/tree/FrequencyTreeDistance.hpp

• View this file on GitHub
• Last update: 2025-11-01 07:26:02+09:00
• Include: #include "library/tree/FrequencyTreeDistance.hpp"

Depends on

• library/fps/fps.hpp
• library/math/modula.hpp
• library/template/template.hpp

Verified with

• verify/yosupo/frequency-tree-distance.test.cpp

Code

#pragma once

#include "../template/template.hpp"
#include "../fps/fps.hpp"
template <class G>
struct FrequenceTreeDistance : CentroidDecomposition<G> {
	using CentroidDecomposition<G>::g;
	using CentroidDecomposition<G>::v;
	using CentroidDecomposition<G>::get_size;
	using CentroidDecomposition<G>::get_centroid;

	FrequenceTreeDistance(const G &_g) : CentroidDecomposition<G>(_g, false) {}

	FPS<ll> count, self;
	void dfs_depth(ll cur, ll par, ll d) {
		while (sz(count) <= d) count.emplace_back(0ll);
		while (sz(self) <= d) self.emplace_back(0ll);
		++count[d];
		++self[d];
		for (ll dst : g[cur]) {
			if (par == dst || v[dst]) continue;
			dfs_depth(dst, cur, d + 1);
		}
	}
	vl get(ll start = 0) {
		count.set_mod(4603910272195756033, 5);
		self.set_mod(4603910272195756033, 5);
		queue<ll> Q;
		Q.push(get_centroid(start, -1, get_size(start, -1) / 2));
		vl ans;
		ans.reserve(sz(g));
		count.reserve(sz(g));
		self.reserve(sz(g));

		while (!Q.empty()) {
			ll r = Q.front();
			Q.pop();
			count.clear();
			v[r] = 1;
			for (auto &c : g[r]) {
				if (v[c]) continue;
				self.clear();
				Q.emplace(get_centroid(c, -1, get_size(c, -1) / 2));
				dfs_depth(c, r, 1);
				auto self2 = self * self;
				while (sz(self2) > sz(ans)) ans.emplace_back(0ll);
				rep(i, sz(self2)) ans[i] -= self2[i];
			}
			if (count.empty()) continue;
			++count[0];
			auto count2 = count * count;
			while (sz(count2) > sz(ans)) ans.emplace_back(0ll);
			rep(i, sz(count2)) ans[i] += count2[i];
		}
		for (auto &x : ans) x >>= 1;
		return ans;
	}
};

#line 2 "library/tree/FrequencyTreeDistance.hpp"

#line 2 "library/template/template.hpp"

#include <bits/stdc++.h>
using namespace std;
#define MM << ' ' <<
using ll = long long;
using ld = long double;
using pll = pair<ll, ll>;
using vl = vector<ll>;
template <class T> using vec = vector<T>;
template <class T> using vv = vec<vec<T>>;
template <class T> using vvv = vv<vec<T>>;
template <class T> using minpq = priority_queue<T, vector<T>, greater<T>>;
#define rep(i, r) for(ll i = 0; i < (r); i++)
#define reps(i, l, r) for(ll i = (l); i < (r); i++)
#define rrep(i, l, r) for(ll i = (r) - 1; i >= l; i--)
template <class T> void uniq(T &a) { sort(all(a)); erase(unique(all(a)), a.end()); }
#define all(a) (a).begin(), (a).end()
#define sz(a) (ll)(a).size()
const ll INF = numeric_limits<ll>::max() / 4;
const ld inf = numeric_limits<ld>::max() / 2;
const ll mod1 = 1000000007;
const ll mod2 = 998244353;
const ld pi = 3.141592653589793238;
template <class T> void rev(T &a) { reverse(all(a)); }
ll popcnt(ll a) { return __builtin_popcountll(a); }
template<typename T>
bool chmax(T &a, const T& b) { return a < b ? a = b, true : false; }
template<typename T>
bool chmin(T &a, const T& b) { return a > b ? a = b, true : false; }
template <typename T1, typename T2>
ostream& operator<<(ostream& os, const pair<T1, T2>& p) {
	os << p.first << " " << p.second;
	return os;
}
template <typename T1, typename T2>
istream& operator>>(istream& is, pair<T1, T2>& p) {
	is >> p.first >> p.second;
	return is;
}
template <typename T>
ostream& operator<<(ostream& os, const vec<T>& v) {
	rep(i, sz(v)) {
		os << v[i] << " \n"[i + 1 == sz(v)];
	}
	return os;
}
template <typename T>
istream& operator>>(istream& is, vec<T>& v) {
	for (T& in : v) is >> in;
	return is;
}
void yesno(bool t) {
	cout << (t ? "Yes" : "No") << endl;
}
#line 2 "library/fps/fps.hpp"

#line 2 "library/math/modula.hpp"

#line 4 "library/math/modula.hpp"

ll modmul(ll a, ll b, ll mod) { return __int128(a) * b % mod; }
ll modpow(ll a, ll b, ll mod) {
  ll ans = 1;
  for (; b; a = modmul(a, a, mod), b /= 2)
    if (b & 1) ans = modmul(ans, a, mod);
  return ans;
}
// a^x == b (mod p), if no such x exists than -1
ll modlog(ll a, ll b, ll p) {
    ll n = (ll)sqrtl(p) + 1, e = 1, f = 1, j = 1;
    unordered_map<ll, ll> A;
    while (j = n && (e = f = e * a % p) != b % p)
        A[e * b % p] = j++;
    if (e == b % p) return j;
    if (gcd(p, e) == gcd(p, b))
        reps(i, 2, n + 2) if (A.count(e = e * f % p)) 
            return n * i - A[e];
    return -1;
}
// x^2 == a (mod p)
ll modsqrt(ll a, ll p) {
    a %= p;
    while (a < 0) a += p;
    if (a == 0) return 0;
    if (modpow(a, (p - 1) / 2, p) != 1) return -1;
    if (p % 4 == 3) return modpow(a, (p + 1) / 4, p);
    ll s = p - 1, n = 2;
    ll r = 0, m;
    while (s % 2 == 0) ++r, s /= 2;
    while (modpow(n, (p - 1) / 2, p) != p - 1) ++n;
    ll x = modpow(a, (s + 1) / 2, p);
    ll b = modpow(a, s, p), g = modpow(n, s, p);
    for (;; r = m) {
        ll t = b;
        for (m = 0; m < r && t != 1; ++m) t = modmul(t, t, p);
        if (m == 0) return x;
        ll gs = modpow(g, 1ll << (r - m - 1), p);
        g = modmul(gs, gs, p);
        x = modmul(x, gs, p);
        b = modmul(b, g, p);
    }
}
ll modinv(ll x, ll mod) { return modpow(x, mod - 2, mod); }
#line 5 "library/fps/fps.hpp"

template <typename T = ll>
struct FPS : vector<T> {
private:
	ll mod = 998244353;
	ll root = 3;
	void ntt(vec<T> &a) {
		ll n = sz(a), L = 31 - __builtin_clz(n);
		vl rt(2, 1);
		for (ll k = 2, s = 2; k < n; k *= 2, s++) {
			rt.resize(n);
			ll z[] = {1, modpow(root, mod >> s, mod)};
			reps(i, k, 2 * k) rt[i] = modmul(rt[i / 2], z[i & 1], mod);
		}
		vec<T> rev(n);
		rep(i, n) rev[i] = (rev[i / 2] | (i & 1) << L) / 2;
		rep(i, n) if (i < rev[i]) swap(a[i], a[rev[i]]);
		for (ll k = 1; k < n; k *= 2) {
			for (ll i = 0; i < n; i += 2 * k) {
				rep(j, k) {
					ll z = modmul(rt[j + k], a[i + j + k], mod), &ai = a[i + j];
					a[i + j + k] = ai - z + (z > ai ? mod : 0ll);
					ai += (ai + z >= mod ? z - mod : z);
				}
			}
		}
	}

	vec<T> conv(vec<T> a, vec<T> b) {
		if (!sz(a) || !sz(b)) return {};
		ll s = sz(a) + sz(b) - 1, B = 32 - __builtin_clz(s), n = 1 << B;
		ll inv = modinv(n, mod);
		vec<T> L(a), R(b), out(n);
		L.resize(n), R.resize(n);
		ntt(L), ntt(R);
		rep(i, n) out[-i & (n - 1)] = modmul(L[i], modmul(R[i], inv, mod), mod);
		ntt(out);
		out.resize(s);
		return out;
	}

public:
	using vector<T>::vector;
	using P = FPS;

	void set_mod(ll md, ll g) {
		mod = md;
		root = g;
	}
	
	P pre(ll def) const {
		return P(begin(*this), begin(*this) + min((ll)this->size(), def));
	}
	P rev(ll deg = -1) const {
		P ret(*this);
		if (deg != -1) ret.resize(deg, T(0));
		reverse(all(ret));
		return ret;
	}

	void shrink() {
		while (this->size() && !this->back()) this->pop_back();
	}

	T freq(ll p) const { return (p < (ll)this->size()) ? (*this)[p] : T(0); }

	P operator+(const P& r) const { return P(*this) += r; }
	P operator+(const T& v) const { return P(*this) += v; }
	P operator-(const P& r) const { return P(*this) -= r; }
	P operator-(const T& v) const { return P(*this) -= v; }
	P operator*(const P& r) const { return P(*this) *= r; }
	P operator*(const T& v) const { return P(*this) *= v; }
	P operator/(const P& r) const { return P(*this) /= r; }
	P operator%(const P& r) const { return P(*this) %= r; }

	P& operator+=(const P& r) {
		if (sz(r) > (ll)this->size()) this->resize(sz(r));
		rep(i, sz(r)) (*this)[i] += r[i];
		return *this;
	}
	P& operator-=(const P& r) {
		if (sz(r) > (ll)this->size()) this->resize(sz(r));
		rep(i, sz(r)) (*this)[i] = (*this)[i] - r[i] + (r[i] > (*this)[i] ? mod : 0ll);
		return *this;
	}

	P& operator*=(const P& r) {
		if (this->empty() || r.empty()) {
			this->clear();
			return *this;
		}
		auto ret = conv(*this, r);
		return *this = {all(ret)};
	}
	P& operator/=(const P& r) {
		if (this->size() < r.size()) {
			this->clear();
			return *this;
		}
		ll n = this->size() - r.size() + 1;
		return *this = (rev().pre(n) * r.rev().inv(n)).pre(n).rev(n);
	}
	P& operator%=(const P& r) {
		*this -= *this / r * r;
		shrink();
		return *this;
	}
	pair<P, P> div_mod(const P& r) {
		P q = *this / r;
		P x = *this - q * r;
		x.shrink();
		return make_pair(q, x);
	}
	P operator-() const {
		P ret((ll)this->size());
		rep(i, (ll)this->size()) ret[i] = -(*this)[i];
		return ret;
	}
	P& operator+=(const T& v) {
		if (this->empty()) this->resize(1);
		(*this)[0] += v;
		return *this;
	}
	P& operator-=(const T& v) {
		if (this->empty()) this->resize(1);
		(*this)[0] -= v;
		return *this;
	}
	P& operator*=(const T& v) {
		rep(i, (ll)this->size()) (*this)[i] = modmul((*this)[i], v, mod);
		return *this;
	}
	P dot(P r) const {
		P ret(min(this->size(), r.size()));
		rep(i, ret.size()) ret[i] = modmul((*this)[i], r[i], mod);
		return ret;
	}
	P operator>>(ll sz) const {
		if ((ll)this->size() <= sz) return {};
		P ret(*this);
		ret.erase(ret.begin(), ret.begin() + sz);
		return ret;
	}
	P operator<<(ll sz) const {
		P ret(*this);
		ret.insert(ret.begin(), sz, T(0));
		return ret;
	}
	T operator()(T x) const {
		T r = 0, w = 1;
		for (auto& v : *this) {
			r += w * v;
			w *= x;
		}
		return r;
	}
	P diff() const {
		const ll n = (ll)this->size();
		P ret(max(0ll, n - 1));
		reps(i, 1, n) ret[i - 1] = modmul((*this)[i], T(i), mod);
		return ret;
	}
	P integral() const {
		const ll n = (ll)this->size();
		P ret(n + 1);
		rep(i, n) ret[i + 1] = modmul(freq(i), modinv(T(i + 1), mod), mod);
		return ret;
	}
	P inv(ll deg = -1) const {
		if (deg == -1) deg = (ll)this->size();
		P res = P({modmul(T(1), modinv(freq(0), mod), mod)});
		for (ll i = 1; i < deg; i *= 2) {
			res = (res * T(2) - res * res * pre(2 * i)).pre(2 * i);
		}
		return res.pre(deg);
	}
	P exp(ll n = -1) const {
		assert(freq(0) == T(0));
		if (n == -1) n = (ll)this->size();
		P g = P({T(1)});
		for (ll i = 1; i < n; i *= 2) {
			g = (g * (pre(i * 2) + P({T(1)}) - g.log(i * 2))).pre(i * 2);
		}
		return g.pre(n);
	}
	P log(ll n = -1) const {
		if (n == -1) n = (ll)this->size();
		assert(freq(0) == T(1));
		auto f = pre(n);
		return (f.diff() * f.inv(n - 1)).pre(n - 1).integral();
	}
	P sqrt(ll deg = -1) const {
		const ll n = (ll)this->size();
		if (deg == -1) deg = n;
		if ((*this)[0] == T(0)) {
			reps(i, 1, n) {
				if ((*this)[i] != T(0)) {
					if (i & 1) return {};
					if (deg - i / 2 <= 0) break;
					auto ret = (*this >> i).sqrt(deg - i / 2);
					if (ret.empty()) return {};
					ret = ret << (i / 2);
					if (sz(ret) < deg) ret.resize(deg, T(0));
					return ret;
				}
			}
			return P(deg);
		}
		auto sqr = T(modsqrt((*this)[0], mod));
		if (abs(modmul(sqr, sqr, mod) - (*this)[0]) % mod != 0) return {};
		P ret{sqr};
		T inv2 = modinv(T(2), mod);
		for (ll i = 1; i < deg; i <<= 1) {
			ret = (ret + pre(i << 1) * ret.inv(i << 1)) * inv2;
		}
		return ret.pre(deg);
	}
	P pow(ll k, ll n = -1) {
		if (n == -1) n = (ll)this->size();
		if (k == 0) {
			P res(n);
			res[0] = T(1);
			return res;
		}
		rep(i, (ll)this->size()) {
			if ((*this)[i]) {
				T rev = modmul(T(1), modinv((*this)[i], mod), mod);
				P ret = (((*this * rev) >> i).log(n) * T(k)).exp(n);
				ret *= modpow((*this)[i], k, mod);
				ret = (ret << (i * k)).pre(n);
				if ((ll)ret.size() < n) ret.resize(n);
				return ret;
			}
			if (__int128(i + 1) * k >= n) return P(n);
		}
		return P(n);
	}
	P pow_mod(ll n, const P& m) {
		P x = *this, r = {{1}};
		while(n) {
			if (n & 1) r = r * x % m;
			x = x * x % m;
			n >>= 1;
		}
		return r;
	}
	P shift(T c) const {
		ll n = this->size();
		vec<T> fact(n), rfact(n);
		fact[0] = rfact[0] = T(1);
		reps(i, 1, n) fact[i] = modmul(fact[i - 1], T(i), mod);
		rfact[n - 1] = T(1) * modinv(fact[n - 1], mod);
		rrep(i, 2, n) rfact[i - 1] = modmul(rfact[i], T(i), mod);
		P p(*this);
		rep(i, n) p[i] = modmul(p[i], fact[i], mod);
		p = p.rev();
		P bs(n, T(1));
		reps(i, 1, n) bs[i] = modmul(bs[i - 1], modmul(c, modmul(rfact[i], fact[i - 1], mod), mod), mod);
		p = (p * bs).pre(n);
		p = p.rev();
		rep(i, n) p[i] = modmul(p[i], rfact[i], mod);
		return p;
	}
};
#line 5 "library/tree/FrequencyTreeDistance.hpp"
template <class G>
struct FrequenceTreeDistance : CentroidDecomposition<G> {
	using CentroidDecomposition<G>::g;
	using CentroidDecomposition<G>::v;
	using CentroidDecomposition<G>::get_size;
	using CentroidDecomposition<G>::get_centroid;

	FrequenceTreeDistance(const G &_g) : CentroidDecomposition<G>(_g, false) {}

	FPS<ll> count, self;
	void dfs_depth(ll cur, ll par, ll d) {
		while (sz(count) <= d) count.emplace_back(0ll);
		while (sz(self) <= d) self.emplace_back(0ll);
		++count[d];
		++self[d];
		for (ll dst : g[cur]) {
			if (par == dst || v[dst]) continue;
			dfs_depth(dst, cur, d + 1);
		}
	}
	vl get(ll start = 0) {
		count.set_mod(4603910272195756033, 5);
		self.set_mod(4603910272195756033, 5);
		queue<ll> Q;
		Q.push(get_centroid(start, -1, get_size(start, -1) / 2));
		vl ans;
		ans.reserve(sz(g));
		count.reserve(sz(g));
		self.reserve(sz(g));

		while (!Q.empty()) {
			ll r = Q.front();
			Q.pop();
			count.clear();
			v[r] = 1;
			for (auto &c : g[r]) {
				if (v[c]) continue;
				self.clear();
				Q.emplace(get_centroid(c, -1, get_size(c, -1) / 2));
				dfs_depth(c, r, 1);
				auto self2 = self * self;
				while (sz(self2) > sz(ans)) ans.emplace_back(0ll);
				rep(i, sz(self2)) ans[i] -= self2[i];
			}
			if (count.empty()) continue;
			++count[0];
			auto count2 = count * count;
			while (sz(count2) > sz(ans)) ans.emplace_back(0ll);
			rep(i, sz(count2)) ans[i] += count2[i];
		}
		for (auto &x : ans) x >>= 1;
		return ans;
	}
};

Back to top page