Skip to content

Instantly share code, notes, and snippets.

@NachiaVivias
Last active March 26, 2024 10:55
Show Gist options
  • Save NachiaVivias/f8a47975913f8627789b904fde1064a6 to your computer and use it in GitHub Desktop.
Save NachiaVivias/f8a47975913f8627789b904fde1064a6 to your computer and use it in GitHub Desktop.
等高線集約クエリのランダムテスト用
// 計算内容:
// 配列 long long A[n][n] が与えられる。はじめ、配列 X[n] の要素はすべて 0 である。
// v = 0..n-1 , d = 0..n-1 について、 v からの距離が d である頂点 x について、 X[x] に A[v][d] を加算する。
// 最終的な配列 X を求めよ。
//#define _GLIBCXX_DEBUG
#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <utility>
#include <cassert>
namespace nachia{
template<class Elem>
class CsrArray{
public:
struct ListRange{
using iterator = typename std::vector<Elem>::iterator;
iterator begi, endi;
iterator begin() const { return begi; }
iterator end() const { return endi; }
int size() const { return (int)std::distance(begi, endi); }
Elem& operator[](int i) const { return begi[i]; }
};
struct ConstListRange{
using iterator = typename std::vector<Elem>::const_iterator;
iterator begi, endi;
iterator begin() const { return begi; }
iterator end() const { return endi; }
int size() const { return (int)std::distance(begi, endi); }
const Elem& operator[](int i) const { return begi[i]; }
};
private:
int m_n;
std::vector<Elem> m_list;
std::vector<int> m_pos;
public:
CsrArray() : m_n(0), m_list(), m_pos() {}
static CsrArray Construct(int n, std::vector<std::pair<int, Elem>> items){
CsrArray res;
res.m_n = n;
std::vector<int> buf(n+1, 0);
for(auto& [u,v] : items){ ++buf[u]; }
for(int i=1; i<=n; i++) buf[i] += buf[i-1];
res.m_list.resize(buf[n]);
for(int i=(int)items.size()-1; i>=0; i--){
res.m_list[--buf[items[i].first]] = std::move(items[i].second);
}
res.m_pos = std::move(buf);
return res;
}
static CsrArray FromRaw(std::vector<Elem> list, std::vector<int> pos){
CsrArray res;
res.m_n = pos.size() - 1;
res.m_list = std::move(list);
res.m_pos = std::move(pos);
return res;
}
ListRange operator[](int u) { return ListRange{ m_list.begin() + m_pos[u], m_list.begin() + m_pos[u+1] }; }
ConstListRange operator[](int u) const { return ConstListRange{ m_list.begin() + m_pos[u], m_list.begin() + m_pos[u+1] }; }
int size() const { return m_n; }
int fullSize() const { return (int)m_list.size(); }
};
} // namespace nachia
namespace nachia{
struct Graph {
public:
struct Edge{
int from, to;
void reverse(){ std::swap(from, to); }
int xorval() const { return from ^ to; }
};
Graph(int n = 0, bool undirected = false, int m = 0) : m_n(n), m_e(m), m_isUndir(undirected) {}
Graph(int n, const std::vector<std::pair<int, int>>& edges, bool undirected = false) : m_n(n), m_isUndir(undirected){
m_e.resize(edges.size());
for(std::size_t i=0; i<edges.size(); i++) m_e[i] = { edges[i].first, edges[i].second };
}
template<class Cin>
static Graph Input(Cin& cin, int n, bool undirected, int m, bool offset = 0){
Graph res(n, undirected, m);
for(int i=0; i<m; i++){
int u, v; cin >> u >> v;
res[i].from = u - offset;
res[i].to = v - offset;
}
return res;
}
int numVertices() const noexcept { return m_n; }
int numEdges() const noexcept { return int(m_e.size()); }
int addNode() noexcept { return m_n++; }
int addEdge(int from, int to){ m_e.push_back({ from, to }); return numEdges() - 1; }
Edge& operator[](int ei) noexcept { return m_e[ei]; }
const Edge& operator[](int ei) const noexcept { return m_e[ei]; }
Edge& at(int ei) { return m_e.at(ei); }
const Edge& at(int ei) const { return m_e.at(ei); }
auto begin(){ return m_e.begin(); }
auto end(){ return m_e.end(); }
auto begin() const { return m_e.begin(); }
auto end() const { return m_e.end(); }
bool isUndirected() const noexcept { return m_isUndir; }
void reverseEdges() noexcept { for(auto& e : m_e) e.reverse(); }
void contract(int newV, const std::vector<int>& mapping){
assert(numVertices() == int(mapping.size()));
for(int i=0; i<numVertices(); i++) assert(0 <= mapping[i] && mapping[i] < newV);
for(auto& e : m_e){ e.from = mapping[e.from]; e.to = mapping[e.to]; }
m_n = newV;
}
std::vector<Graph> induce(int num, const std::vector<int>& mapping) const {
int n = numVertices();
assert(n == int(mapping.size()));
for(int i=0; i<n; i++) assert(-1 <= mapping[i] && mapping[i] < num);
std::vector<int> indexV(n), newV(num);
for(int i=0; i<n; i++) if(mapping[i] >= 0) indexV[i] = newV[mapping[i]]++;
std::vector<Graph> res; res.reserve(num);
for(int i=0; i<num; i++) res.emplace_back(newV[i], isUndirected());
for(auto e : m_e) if(mapping[e.from] == mapping[e.to] && mapping[e.to] >= 0) res[mapping[e.to]].addEdge(indexV[e.from], indexV[e.to]);
return res;
}
CsrArray<int> getEdgeIndexArray(bool undirected) const {
std::vector<std::pair<int, int>> src;
src.reserve(numEdges() * (undirected ? 2 : 1));
for(int i=0; i<numEdges(); i++){
auto e = operator[](i);
src.emplace_back(e.from, i);
if(undirected) src.emplace_back(e.to, i);
}
return CsrArray<int>::Construct(numVertices(), src);
}
CsrArray<int> getEdgeIndexArray() const { return getEdgeIndexArray(isUndirected()); }
CsrArray<int> getAdjacencyArray(bool undirected) const {
std::vector<std::pair<int, int>> src;
src.reserve(numEdges() * (undirected ? 2 : 1));
for(auto e : m_e){
src.emplace_back(e.from, e.to);
if(undirected) src.emplace_back(e.to, e.from);
}
return CsrArray<int>::Construct(numVertices(), src);
}
CsrArray<int> getAdjacencyArray() const { return getAdjacencyArray(isUndirected()); }
private:
int m_n;
std::vector<Edge> m_e;
bool m_isUndir;
};
} // namespace nachia
namespace nachia {
struct TreeHorizontalSets {
int n;
int numHorizontalSets;
int emptySet;
std::vector<int> baseDepth;
std::vector<int> descend;
std::vector<int> expand;
std::vector<int> ascend;
TreeHorizontalSets(const Graph& tree, int root = 0)
: n(tree.numVertices())
{
int n = tree.numVertices();
auto adj = tree.getAdjacencyArray();
emptySet = n*2-1;
std::vector<int> bfs(n, -1);
std::vector<int> par(n, -1);
baseDepth.assign(n*2, 0);
bfs[0] = root;
int bfsi = 1;
for(int i=0; i<n; i++){
int v = bfs[i];
int eiz = adj[v].size();
for(int ei=0; ei<eiz; ){
int w = adj[v][ei];
if(w == par[v]){
std::swap(adj[v][--eiz], adj[v][ei]);
continue;
}
par[w] = v;
baseDepth[w] = baseDepth[v] + 1;
bfs[bfsi++] = w;
ei++;
}
}
int treeHeight = baseDepth[bfs[n-1]];
int nextNodeId = n;
descend.assign(n*2, emptySet);
expand.assign(n*2, emptySet);
ascend.assign(n*2, emptySet);
std::vector<int> height(n, 0);
for(int i=n-1; i>=0; i--){
int v = bfs[i];
int h1 = 0, h2 = 0, h1w = emptySet;
int depthv = baseDepth[v];
for(int w : adj[v]) if(par[v] != w){
int h = height[w];
if(h1 < h){ std::swap(h1, h); h1w = w; }
if(h2 < h) std::swap(h2, h);
}
{
int p = v, q = h1w;
for(int h=0; h<h2; h++){
int nx = nextNodeId++;
ascend[nx] = p;
baseDepth[nx] = depthv;
descend[p] = nx;
p = nx; q = descend[q];
}
descend[p] = q;
}
for(int w : adj[v]) if(par[v] != w){
ascend[w] = v;
int p = descend[v], q = w, h = std::min(height[w], h2);
while(h--){
expand[q] = p;
p = descend[p]; q = descend[q];
}
}
height[v] = h1 + 1;
}
numHorizontalSets = nextNodeId;
}
};
} // namespace nachia
namespace nachia{
struct HeavyLightDecomposition{
private:
int N;
std::vector<int> P;
std::vector<int> PP;
std::vector<int> PD;
std::vector<int> D;
std::vector<int> I;
std::vector<int> rangeL;
std::vector<int> rangeR;
public:
HeavyLightDecomposition(const CsrArray<int>& E = CsrArray<int>::Construct(1, {}), int root = 0){
N = E.size();
P.assign(N, -1);
I.assign(N, 0); I[0] = root;
int iI = 1;
for(int i=0; i<iI; i++){
int p = I[i];
for(int e : E[p]) if(P[p] != e){
I[iI++] = e;
P[e] = p;
}
}
std::vector<int> Z(N, 1);
std::vector<int> nx(N, -1);
PP.resize(N);
for(int i=0; i<N; i++) PP[i] = i;
for(int i=N-1; i>=1; i--){
int p = I[i];
Z[P[p]] += Z[p];
if(nx[P[p]] == -1) nx[P[p]] = p;
if(Z[nx[P[p]]] < Z[p]) nx[P[p]] = p;
}
for(int p : I) if(nx[p] != -1) PP[nx[p]] = p;
PD.assign(N,N);
PD[root] = 0;
D.assign(N,0);
for(int p : I) if(p != root){
PP[p] = PP[PP[p]];
PD[p] = std::min(PD[PP[p]], PD[P[p]]+1);
D[p] = D[P[p]]+1;
}
rangeL.assign(N,0);
rangeR.assign(N,0);
for(int p : I){
rangeR[p] = rangeL[p] + Z[p];
int ir = rangeR[p];
for(int e : E[p]) if(P[p] != e) if(e != nx[p]){
rangeL[e] = (ir -= Z[e]);
}
if(nx[p] != -1){
rangeL[nx[p]] = rangeL[p] + 1;
}
}
for(int i=0; i<N; i++) I[rangeL[i]] = i;
}
HeavyLightDecomposition(const Graph& tree, int root = 0)
: HeavyLightDecomposition(tree.getAdjacencyArray(true), root) {}
int numVertices() const { return N; }
int depth(int p) const { return D[p]; }
int toSeq(int vtx) const { return rangeL[vtx]; }
int toVtx(int seqidx) const { return I[seqidx]; }
int toSeq2In(int vtx) const { return rangeL[vtx] * 2 - D[vtx]; }
int toSeq2Out(int vtx) const { return rangeR[vtx] * 2 - D[vtx] - 1; }
int parentOf(int v) const { return P[v]; }
int heavyRootOf(int v) const { return PP[v]; }
int heavyChildOf(int v) const {
if(toSeq(v) == N-1) return -1;
int cand = toVtx(toSeq(v) + 1);
if(PP[v] == PP[cand]) return cand;
return -1;
}
int lca(int u, int v) const {
if(PD[u] < PD[v]) std::swap(u, v);
while(PD[u] > PD[v]) u = P[PP[u]];
while(PP[u] != PP[v]){ u = P[PP[u]]; v = P[PP[v]]; }
return (D[u] > D[v]) ? v : u;
}
int dist(int u, int v) const {
return depth(u) + depth(v) - depth(lca(u,v)) * 2;
}
struct Range{
int l; int r;
int size() const { return r-l; }
bool includes(int x) const { return l <= x && x < r; }
};
std::vector<Range> path(int r, int c, bool include_root = true, bool reverse_path = false) const {
if(PD[c] < PD[r]) return {};
std::vector<Range> res(PD[c]-PD[r]+1);
for(int i=0; i<(int)res.size()-1; i++){
res[i] = { rangeL[PP[c]], rangeL[c]+1 };
c = P[PP[c]];
}
if(PP[r] != PP[c] || D[r] > D[c]) return {};
res.back() = { rangeL[r]+(include_root?0:1), rangeL[c]+1 };
if(res.back().l == res.back().r) res.pop_back();
if(!reverse_path) std::reverse(res.begin(),res.end());
else for(auto& a : res) a = { N - a.r, N - a.l };
return res;
}
Range subtree(int p) const { return { rangeL[p], rangeR[p] }; }
int median(int x, int y, int z) const {
return lca(x,y) ^ lca(y,z) ^ lca(x,z);
}
int la(int from, int to, int d) const {
if(d < 0) return -1;
int g = lca(from,to);
int dist0 = D[from] - D[g] * 2 + D[to];
if(dist0 < d) return -1;
int p = from;
if(D[from] - D[g] < d){ p = to; d = dist0 - d; }
while(D[p] - D[PP[p]] < d){
d -= D[p] - D[PP[p]] + 1;
p = P[PP[p]];
}
return I[rangeL[p] - d];
}
struct ChildrenIterRange {
struct Iter {
const HeavyLightDecomposition& hld; int s;
int operator*() const { return hld.toVtx(s); }
Iter& operator++(){
s += hld.subtree(hld.I[s]).size();
return *this;
}
Iter operator++(int) const { auto a = *this; return ++a; }
bool operator==(Iter& r) const { return s == r.s; }
bool operator!=(Iter& r) const { return s != r.s; }
};
const HeavyLightDecomposition& hld; int v;
Iter begin() const { return { hld, hld.rangeL[v] + 1 }; }
Iter end() const { return { hld, hld.rangeR[v] }; }
};
ChildrenIterRange children(int v) const {
return ChildrenIterRange{ *this, v };
}
};
} // namespace nachia
#include <random>
#include <chrono>
#include <unordered_map>
#include <cstdint>
#include <array>
namespace nachia{
class Xoshiro256pp{
public:
using i32 = int32_t;
using u32 = uint32_t;
using i64 = int64_t;
using u64 = uint64_t;
private:
std::array<u64, 4> s;
// https://prng.di.unimi.it/xoshiro256plusplus.c
static inline uint64_t rotl(const uint64_t x, int k) noexcept {
return (x << k) | (x >> (64 - k));
}
inline uint64_t gen(void) noexcept {
const uint64_t result = rotl(s[0] + s[3], 23) + s[0];
const uint64_t t = s[1] << 17;
s[2] ^= s[0];
s[3] ^= s[1];
s[1] ^= s[2];
s[0] ^= s[3];
s[2] ^= t;
s[3] = rotl(s[3], 45);
return result;
}
// https://xoshiro.di.unimi.it/splitmix64.c
u64 splitmix64(u64& x) {
u64 z = (x += 0x9e3779b97f4a7c15);
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9;
z = (z ^ (z >> 27)) * 0x94d049bb133111eb;
return z ^ (z >> 31);
}
public:
void seed(u64 x = 7001){
assert(x != 0);
s[0] = x;
for(int i=1; i<4; i++) s[i] = splitmix64(x);
}
std::array<u64, 4> getState() const { return s; }
void setState(std::array<u64, 4> a){ s = a; }
Xoshiro256pp(){ seed(); }
u64 rng64() { return gen(); }
u64 operator()(){ return gen(); }
// generate x : l <= x <= r
u64 random_unsigned(u64 l,u64 r){
assert(l<=r);
r-=l;
auto res = rng64();
if(res<=r) return res+l;
u64 d = r+1;
u64 max_valid = 0xffffffffffffffff/d*d;
while(true){
auto res = rng64();
if(res<=max_valid) break;
}
return res%d+l;
}
// generate x : l <= x <= r
i64 random_signed(i64 l,i64 r){
assert(l<=r);
u64 unsigned_l = (u64)l ^ (1ull<<63);
u64 unsigned_r = (u64)r ^ (1ull<<63);
u64 unsigned_res = random_unsigned(unsigned_l,unsigned_r) ^ (1ull<<63);
return (i64)unsigned_res;
}
// permute x : n_left <= x <= n_right
// output r from the front
template<class Int>
std::vector<Int> random_nPr(Int n_left, Int n_right, Int r){
Int n = n_right-n_left;
assert(n>=0);
assert(r<=(1ll<<27));
if(r==0) return {};
assert(n>=r-1);
std::vector<Int> V;
std::unordered_map<Int,Int> G;
for(int i=0; i<r; i++){
Int p = random_signed(i,n);
Int x = p - G[p];
V.push_back(x);
G[p] = p - (i - G[i]);
}
for(Int& v : V) v+=n_left;
return V;
}
// V[i] := V[perm[i]]
// using swap
template<class E,class PermInt_t>
void permute_inplace(std::vector<E>& V,std::vector<PermInt_t> perm){
assert(V.size() == perm.size());
int N=V.size();
for(int i=0; i<N; i++){
int p=i;
while(perm[p]!=i){
assert(0 <= perm[p] && perm[p] < N);
assert(perm[p] != perm[perm[p]]);
std::swap(V[p],V[perm[p]]);
int pbuf = perm[p]; perm[p] = p; p = pbuf;
}
perm[p] = p;
}
}
template<class E>
std::vector<E> shuffle(const std::vector<E>& V){
int N=V.size();
auto P = random_nPr(0,N-1,N);
std::vector<E> res;
res.reserve(N);
for(int i=0; i<N; i++) res.push_back(V[P[i]]);
return res;
}
// shuffle using swap
template<class E>
void shuffle_inplace(std::vector<E>& V){
int N=V.size();
permute_inplace(V,random_nPr(0,N-1,N));
}
};
} // namespace nachia
nachia::Xoshiro256pp rng;
namespace RngInitInstance { struct RngInit { RngInit(){
unsigned long long seed1 = std::random_device()();
unsigned long long seed2 = std::chrono::high_resolution_clock::now().time_since_epoch().count();
rng.seed(seed1 ^ seed2);
} } a; }
using i64 = long long;
#define rep(i,n) for(int i=0; i<int(n); i++)
#define repr(i,n) for(int i=int(n)-1; i>=0; i--)
template<typename A> void chmin(A& l, const A& r){ if(r < l) l = r; }
template<typename A> void chmax(A& l, const A& r){ if(l < r) l = r; }
using namespace std;
vector<i64> fast(int N, const nachia::Graph& tree, const vector<vector<i64>>& A){
int root = 0;
auto adj = tree.getAdjacencyArray();
auto treeHld = nachia::HeavyLightDecomposition(tree, root);
auto hor = nachia::TreeHorizontalSets(tree, root);
int n2 = hor.numHorizontalSets;
auto decTree = nachia::Graph(n2 + 1, true);
rep(v,n2){
int w = hor.descend[v];
if(w == hor.emptySet) w = n2;
decTree.addEdge(v, w);
}
vector<int> heightVariant(N);
repr(i,N) if(i){
int v = treeHld.toVtx(i);
chmax(heightVariant[treeHld.parentOf(v)], heightVariant[v] + 1);
}
rep(i,N) heightVariant[i] -= treeHld.depth(i);
auto decHld = nachia::HeavyLightDecomposition(decTree, n2);
vector<int> remLink(n2+1, n2);
vector<int> linkOrder;
vector<int> fromDq(N*2, n2); // offset +N
vector<vector<pair<int,int>>> nearRem(N*2);
auto dfs = [&](auto& dfs, int v) -> void {
int depth_v = treeHld.depth(v);
int X = v;
int depth_f = depth_v;
while(X != hor.emptySet){
int exX = hor.expand[X];
if(exX == hor.emptySet || hor.baseDepth[exX] != depth_v - 1) break;
int dq = depth_f - hor.baseDepth[exX] * 2;
remLink[X] = fromDq[N+dq];
linkOrder.push_back(X);
fromDq[N+dq] = X;
nearRem[N+dq].push_back(make_pair(treeHld.subtree(v).l, X));
X = hor.descend[X]; depth_f++;
}
for(int w : treeHld.children(v)) dfs(dfs, w);
X = v;
depth_f = depth_v;
while(X != hor.emptySet){
int exX = hor.expand[X];
if(exX == hor.emptySet || hor.baseDepth[exX] != depth_v - 1) break;
int dq = depth_f - hor.baseDepth[exX] * 2;
fromDq[N+dq] = remLink[X];
nearRem[N+dq].push_back(make_pair(treeHld.subtree(v).r, remLink[X]));
X = hor.descend[X]; depth_f++;
}
};
dfs(dfs, root);
auto preprocess = [&](int dq, int v) -> int {
if(heightVariant[v] >= dq - 2) return v;
if(!(heightVariant[root] >= dq - 2)) return -1;
while(true){
int w = treeHld.heavyRootOf(v);
if(w < 0 || heightVariant[w] >= dq - 2) break;
v = treeHld.parentOf(w);
}
int wi = treeHld.toSeq(treeHld.heavyRootOf(v));
int vi = treeHld.toSeq(v) + 1;
while(wi + 1 < vi){
int x = (wi + vi) / 2;
bool ok = (heightVariant[treeHld.toVtx(x)] >= dq - 2);
(ok ? wi : vi) = x;
}
return treeHld.toVtx(wi);
};
auto searchNearRem = [&](int dq, int v) -> int {
auto& tg = nearRem[dq+N];
auto i = upper_bound(tg.begin(), tg.end(), make_pair(treeHld.toSeq(v), n2+1)) - tg.begin();
if(i == 0) return n2;
return tg[i-1].second;
};
vector<i64> buf1(n2+1);
vector<i64> buf2(n2+1);
rep(vraw,N) rep(draw,N){
i64 a = A[vraw][draw];
int dq = draw - treeHld.depth(vraw);
int v = preprocess(dq, vraw);
if(v < 0) continue;
int d = draw - (treeHld.depth(vraw) - treeHld.depth(v));
if(decHld.depth(v) > d){
buf2[decHld.la(v, n2, d)] += a;
}
if(d != 0){
if(d != 0 && treeHld.depth(v) >= d){
buf2[treeHld.la(v, root, d)] += a;
}
int near = searchNearRem(dq, v);
if(near != n2) buf1[near] += a;
}
}
reverse(linkOrder.begin(), linkOrder.end());
for(int x : linkOrder) buf1[remLink[x]] += buf1[x];
rep(i,n2) if(hor.expand[i] != hor.emptySet){
buf2[i] -= buf1[i];
buf2[hor.expand[i]] += buf1[i];
}
repr(i,n2) if(hor.expand[i] != hor.emptySet) buf2[i] += buf2[hor.expand[i]];
vector<i64> ans(N);
rep(i,N) ans[i] = buf2[i];
return ans;
}
vector<i64> naive(int N, const nachia::Graph& tree, const vector<vector<i64>>& A){
auto hld = nachia::HeavyLightDecomposition(tree);
vector<i64> ans(N);
rep(s,N) rep(t,N) ans[t] += A[s][hld.dist(s, t)];
return ans;
}
void testcase(int caseid){
auto rngState = rng.getState();
cout << "testcase id #" << caseid << endl;
int N = 1000;
bool doRandomTree = true;
vector<vector<i64>> A(N, vector<i64>(N));
rep(i,N) rep(j,N) A[i][j] = rng.random_signed(1, 1000000000);
auto tree = nachia::Graph(N, true);
if(doRandomTree){
vector<int> P(N); rep(i,N) P[i] = i;
rng.shuffle_inplace(P);
rep(i,N-1) tree.addEdge(P[rng.random_signed(max(0,i-10),i)], P[i+1]);
} else {
rep(i,N-1){
int u,v; cin >> u >> v; tree.addEdge(u,v);
}
}
auto ans_naive = naive(N, tree, A);
auto ans_fast = fast(N, tree, A);
if(ans_naive != ans_fast){
cout << "wrong answer" << endl;
cout << "testcase id #" << caseid << endl;
cout << "rngState = {";
rep(i,rngState.size()){ if(i){ cout << ","; } cout << rngState[i]; }
cout << endl;
cout << " --------- " << endl;
cout << N << endl;
for(auto [u,v] : tree) cout << u << " " << v << endl;
exit(0);
}
}
int main(){
ios::sync_with_stdio(false); cin.tie(nullptr);
int T = 30;
for(int t=0; t<T; ++t) testcase(t);
cout << "all verified" << endl;
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment