【LCA DFS 前缀和】P10391 [蓝桥杯 2024 省 A] 零食采购|普及+

本文涉及知识点

C++算法：前缀和、前缀乘积、前缀异或的原理、源码及测试用例 包括课程视频
C++DFS
倍增算法(multiply)、树上倍增、最近公共祖先(LCA)

P10391 [蓝桥杯 2024 省 A] 零食采购

题目描述

小蓝准备去星际旅行，出发前想在本星系采购一些零食，星系内有 n n n 颗星球，由 n − 1 n-1 n−1 条航路连接为连通图，第 i i i 颗星球卖第 c i c_i ci​ 种零食特产。小蓝想出了 q q q 个采购方案，第 i i i 个方案的起点为星球 s i s_i si​ ，终点为星球 t i t_i ti​ ，对于每种采购方案，小蓝将从起点走最短的航路到终点，并且可以购买所有经过的星球上的零食（包括起点终点），请计算每种采购方案最多能买多少种不同的零食。

输入格式

输入的第一行包含两个正整数 n n n， q q q，用一个空格分隔。
第二行包含 n n n 个整数 c 1 , c 2 , ⋯ , c n c_1,c_2,\cdots, c_n c1​,c2​,⋯,cn​，相邻整数之间使用一个空格分隔。
接下来 n − 1 n - 1 n−1 行，第 i i i 行包含两个整数 u i , v i u_i,v_i ui​,vi​，用一个空格分隔，表示一条
航路将星球 u i u_i ui​ 与 v i v_i vi​ 相连。
接下来 q q q 行，第 i i i 行包含两个整数 $s_i
, t_i $，用一个空格分隔，表示一个采购方案。

输出格式

输出 q q q 行，每行包含一个整数，依次表示每个采购方案的答案。

输入输出样例 #1

输入 #1

4 2 1 2 3 1 1 2 1 3 2 4 4 3 1 4 

输出 #1

3 2 

说明/提示

第一个方案路线为 { 4 , 2 , 1 , 3 } \{4, 2, 1, 3\} {4,2,1,3}，可以买到第 1 , 2 , 3 1, 2, 3 1,2,3 种零食；
第二个方案路线为 { 1 , 2 , 4 } \{1, 2, 4\} {1,2,4}，可以买到第 1 , 2 1, 2 1,2 种零食。

对于 20% 的评测用例，$1 ≤ n, q ≤ 5000 $；
对于所有评测用例， 1 ≤ n , q ≤ 1 0 5 ， 1 ≤ c i ≤ 20 ， 1 ≤ u i , v i ≤ n ， 1 ≤ s i , t i ≤ n 1 ≤ n, q ≤ 10^5，1 ≤ c_i ≤ 20，1 ≤ u_i , v_i ≤ n，1 ≤ s_i , t_i ≤ n 1≤n,q≤105，1≤ci​≤20，1≤ui​,vi​≤n，1≤si​,ti​≤n。

DFS 树上前缀和 LCA

以1(0)为根，cnt[i][j]记录第j个星球是否有货物i。preSum[i][j]，节点j到根节点整个路径包括货物i的星球数量。初始化只需要一次DFS。时间复杂度：O(20n)
每次查询，令u和v的最近公共祖先g，如果preSum[i][u]+preSum[i][v]-2preSum[i][g]+cnt[i][g] > 0，则可以买到货物i。
是否复杂度：O(20qlogn)

代码

核心代码

#include<iostream>#include<sstream>#include<vector>#include<map>#include<unordered_map>#include<set>#include<unordered_set>#include<string>#include<algorithm>#include<functional>#include<queue>#include<stack>#include<iomanip>#include<numeric>#include<math.h>#include<climits>#include<assert.h>#include<cstring>#include<list>#include<array>#include<bitset>usingnamespace std;template<classT1,classT2> std::istream&operator>>(std::istream& in, pair<T1, T2>& pr){ in >> pr.first >> pr.second;return in;}template<classT1,classT2,classT3> std::istream&operator>>(std::istream& in, tuple<T1, T2, T3>& t){ in >> get<0>(t)>> get<1>(t)>> get<2>(t);return in;}template<classT1,classT2,classT3,classT4> std::istream&operator>>(std::istream& in, tuple<T1, T2, T3, T4>& t){ in >> get<0>(t)>> get<1>(t)>> get<2>(t)>> get<3>(t);return in;}template<classT1,classT2,classT3,classT4,classT5,classT6,classT7> std::istream&operator>>(std::istream& in, tuple<T1, T2, T3, T4,T5,T6,T7>& t){ in >> get<0>(t)>> get<1>(t)>> get<2>(t)>> get<3>(t)>> get<4>(t)>> get<5>(t)>> get<6>(t);return in;}template<classT=int> vector<T>Read(){int n; cin >> n; vector<T>ret(n);for(int i =0; i < n; i++){ cin >> ret[i];}return ret;}template<classT=int> vector<T>ReadNotNum(){ vector<T> ret; T tmp;while(cin >> tmp){ ret.emplace_back(tmp);if('\n'== cin.get()){break;}}return ret;}template<classT=int> vector<T>Read(int n){ vector<T>ret(n);for(int i =0; i < n; i++){ cin >> ret[i];}return ret;}template<int N =1'000'000>classCOutBuff{public:COutBuff(){ m_p = puffer;}template<classT>voidwrite(T x){int num[28], sp =0;if(x <0)*m_p++='-', x =-x;if(!x)*m_p++=48;while(x) num[++sp]= x %10, x /=10;while(sp)*m_p++= num[sp--]+48;AuotToFile();}voidwritestr(constchar* sz){strcpy(m_p, sz); m_p +=strlen(sz);AuotToFile();}inlinevoidwrite(char ch){*m_p++= ch;AuotToFile();}inlinevoidToFile(){fwrite(puffer,1, m_p - puffer,stdout); m_p = puffer;}~COutBuff(){ToFile();}private:inlinevoidAuotToFile(){if(m_p - puffer > N -100){ToFile();}}char puffer[N],* m_p;};template<int N =1'000'000>classCInBuff{public:inlineCInBuff(){}inline CInBuff<N>&operator>>(char& ch){FileToBuf();while(('\r'==*S)||('\n'==*S)||(' '==*S)){ S++;}//忽略空格和回车 ch =*S++;return*this;}inline CInBuff<N>&operator>>(int& val){FileToBuf();intx(0),f(0);while(!isdigit(*S)) f |=(*S++=='-');while(isdigit(*S)) x =(x <<1)+(x <<3)+(*S++^48); val = f ?-x : x; S++;//忽略空格换行 return*this;}inline CInBuff&operator>>(longlong& val){FileToBuf();longlongx(0);intf(0);while(!isdigit(*S)) f |=(*S++=='-');while(isdigit(*S)) x =(x <<1)+(x <<3)+(*S++^48); val = f ?-x : x; S++;//忽略空格换行return*this;}template<classT1,classT2>inline CInBuff&operator>>(pair<T1, T2>& val){*this>> val.first >> val.second;return*this;}template<classT1,classT2,classT3>inline CInBuff&operator>>(tuple<T1, T2, T3>& val){*this>> get<0>(val)>> get<1>(val)>> get<2>(val);return*this;}template<classT1,classT2,classT3,classT4>inline CInBuff&operator>>(tuple<T1, T2, T3, T4>& val){*this>> get<0>(val)>> get<1>(val)>> get<2>(val)>> get<3>(val);return*this;}template<classT=int>inline CInBuff&operator>>(vector<T>& val){int n;*this>> n; val.resize(n);for(int i =0; i < n; i++){*this>> val[i];}return*this;}template<classT=int> vector<T>Read(int n){ vector<T>ret(n);for(int i =0; i < n; i++){*this>> ret[i];}return ret;}template<classT=int> vector<T>Read(){ vector<T> ret;*this>> ret;return ret;}private:inlinevoidFileToBuf(){constint canRead = m_iWritePos -(S - buffer);if(canRead >=100){return;}if(m_bFinish){return;}for(int i =0; i < canRead; i++){ buffer[i]= S[i];//memcpy出错 } m_iWritePos = canRead; buffer[m_iWritePos]=0; S = buffer;int readCnt =fread(buffer + m_iWritePos,1, N - m_iWritePos,stdin);if(readCnt <=0){ m_bFinish =true;return;} m_iWritePos += readCnt; buffer[m_iWritePos]=0; S = buffer;}int m_iWritePos =0;bool m_bFinish =false;char buffer[N +10],* S = buffer;};classCNeiBo{public:static vector<vector<int>>Two(int n,const vector<pair<int,int>>& edges,bool bDirect,int iBase =0){ vector<vector<int>>vNeiBo(n);for(constauto&[i1, i2]: edges){ vNeiBo[i1 - iBase].emplace_back(i2 - iBase);if(!bDirect){ vNeiBo[i2 - iBase].emplace_back(i1 - iBase);}}return vNeiBo;}static vector<vector<int>>Two(int n,const vector<vector<int>>& edges,bool bDirect,int iBase =0){ vector<vector<int>>vNeiBo(n);for(constauto& v : edges){ vNeiBo[v[0]- iBase].emplace_back(v[1]- iBase);if(!bDirect){ vNeiBo[v[1]- iBase].emplace_back(v[0]- iBase);}}return vNeiBo;}static vector<vector<std::pair<int,int>>>Three(int n, vector<vector<int>>& edges,bool bDirect,int iBase =0){ vector<vector<std::pair<int,int>>>vNeiBo(n);for(constauto& v : edges){ vNeiBo[v[0]- iBase].emplace_back(v[1]- iBase, v[2]);if(!bDirect){ vNeiBo[v[1]- iBase].emplace_back(v[0]- iBase, v[2]);}}return vNeiBo;}static vector<vector<std::pair<int,int>>>Three(int n,const vector<tuple<int,int,int>>& edges,bool bDirect,int iBase =0){ vector<vector<std::pair<int,int>>>vNeiBo(n);for(constauto&[u, v, w]: edges){ vNeiBo[u - iBase].emplace_back(v - iBase, w);if(!bDirect){ vNeiBo[v - iBase].emplace_back(u - iBase, w);}}return vNeiBo;}static vector<vector<int>>Mat(vector<vector<int>>& neiBoMat){ vector<vector<int>>neiBo(neiBoMat.size());for(int i =0; i < neiBoMat.size(); i++){for(int j = i +1; j < neiBoMat.size(); j++){if(neiBoMat[i][j]){ neiBo[i].emplace_back(j); neiBo[j].emplace_back(i);}}}return neiBo;}};classCBFSLeve{public:static vector<int>Leve(const vector<vector<int>>& neiBo, vector<int> start){ vector<int>leves(neiBo.size(),-1);for(constauto& s : start){ leves[s]=0;}for(int i =0; i < start.size(); i++){for(constauto& next : neiBo[start[i]]){if(-1!= leves[next]){continue;} leves[next]= leves[start[i]]+1; start.emplace_back(next);}}return leves;}template<classNextFun>static vector<int>Leve(int N, NextFun nextFun, vector<int> start){ vector<int>leves(N,-1);for(constauto& s : start){ leves[s]=0;}for(int i =0; i < start.size(); i++){auto nexts =nextFun(start[i]);for(constauto& next : nexts){if(-1!= leves[next]){continue;} leves[next]= leves[start[i]]+1; start.emplace_back(next);}}return leves;}static vector<vector<int>>LeveNodes(const vector<int>& leves){constint iMaxLeve =*max_element(leves.begin(), leves.end()); vector<vector<int>>ret(iMaxLeve +1);for(int i =0; i < leves.size(); i++){ ret[leves[i]].emplace_back(i);}return ret;};static vector<int>LeveSort(const vector<int>& leves){constint iMaxLeve =*max_element(leves.begin(), leves.end()); vector<vector<int>>leveNodes(iMaxLeve +1);for(int i =0; i < leves.size(); i++){ leveNodes[leves[i]].emplace_back(i);} vector<int> ret;for(constauto& v : leveNodes){ ret.insert(ret.end(), v.begin(), v.end());}return ret;};};classCParents{public:CParents(vector<int>& vParent,longlong iMaxDepth){int iBitNum =0;for(; iMaxDepth; iBitNum++){constauto mask =1LL<< iBitNum;if(mask & iMaxDepth){ iMaxDepth = iMaxDepth ^ mask;}}constint n = vParent.size(); m_vParents.assign(iBitNum +1, vector<int>(n,-1)); m_vParents[0]= vParent;//树上倍增for(int i =1; i < m_vParents.size(); i++){for(int j =0; j < n; j++){constint iPre = m_vParents[i -1][j];if(-1!= iPre){ m_vParents[i][j]= m_vParents[i -1][iPre];}}}}intGetParent(int iNode,int iDepth)const{int iParent = iNode;for(int iBit =0; iBit < m_vParents.size(); iBit++){if(-1== iParent){return iParent;}if(iDepth &(1<< iBit)){ iParent = m_vParents[iBit][iParent];}}return iParent;}inlineintGetBitCnt()const{return m_vParents.size();};inlineconstint&GetPow2Parent(int iNode,int n)const{return m_vParents[n][iNode];}//在leftNodeExclude的1到rightLeve级祖先中查找符合pr的最近祖先template<class_Pr>intFindFirst(int leftNodeExclude,int rightLeve, _Pr pr){for(int iBit =GetBitCnt()-1; iBit >=0; iBit--){constint iMask =1<< iBit;if(!(iMask & rightLeve)){continue;}if(pr(m_vParents[iBit][leftNodeExclude])){returnBFindFirst(leftNodeExclude, iBit, pr);} leftNodeExclude = m_vParents[iBit][leftNodeExclude];}return-1;}//在node的0到rightLeve级祖先中查找符合pr的最远祖先比node高多少层次,这些层次必须存在template<class_Pr>intFindEnd(int node,int rightLeve, _Pr pr){int leve =0;for(int iBit =GetBitCnt()-1; iBit >=0; iBit--){constint iMask =1<< iBit;if(!(iMask & rightLeve)){continue;}if(!pr(m_vParents[iBit][node])){return leve +BFindEnd(node, iBit, pr);} node = m_vParents[iBit][node]; leve = leve ^ iMask;}return leve;}protected://在leftNodeExclude的1到2^pow^级祖先中查找符合pr的最近祖先template<class_Pr>intBFindFirst(int leftNodeExclude,int pow, _Pr pr){while(pow >0){constint& mid = m_vParents[pow -1][leftNodeExclude];if(pr(mid)){}else{ leftNodeExclude = mid;} pow--;}return m_vParents[0][leftNodeExclude];}//在node的[0,2^pow^-1]级祖先中寻找符合的最后一个template<class_Pr>intBFindEnd(int node,int pow, _Pr pr){int leve =0;while(pow >0){ pow--;constint& mid = m_vParents[pow][node];if(pr(mid)){ node = mid; leve = leve ^(1<< pow);}else{}}return leve;} vector<vector<int>> m_vParents;};classC2Parents:publicCParents{public:C2Parents(vector<int>& vParent,const vector<int>& vDepth):m_vDepth(vDepth),CParents(vParent,*std::max_element(vDepth.begin(), vDepth.end())){}intGetPublicParent(int iNode1,int iNode2)const{int leve0 = m_vDepth[iNode1];int leve1 = m_vDepth[iNode2];if(leve0 < leve1){ iNode2 =GetParent(iNode2, leve1 - leve0); leve1 = leve0;}else{ iNode1 =GetParent(iNode1, leve0 - leve1); leve0 = leve1;}if(iNode1 == iNode2){return iNode1;}for(int iBit =GetBitCnt()-1; iBit >=0; iBit--){constint iMask =1<< iBit;if(iMask & leve0){constint i1 =GetPow2Parent(iNode1, iBit);constint i2 =GetPow2Parent(iNode2, iBit);if(i1 == i2){while(iBit >0){constint i3 =GetPow2Parent(iNode1, iBit -1);constint i4 =GetPow2Parent(iNode2, iBit -1);if(i3 != i4){ iNode1 = i3; iNode2 = i4;} iBit--;}returnGetPow2Parent(iNode1,0);}else{ iNode1 = i1; iNode2 = i2; leve0 -= iMask;}}}return iNode1;}protected: vector<vector<int>> m_vParents;const vector<int> m_vDepth;};classSolution{public: vector<int>Ans(constint N, vector<int>& c, vector<pair<int,int>>& edge, vector<pair<int,int>>& que){auto neiBo =CNeiBo::Two(N, edge,false,1); vector<vector<int>>cnt(20, vector<int>(N)),preSum(20, vector<int>(N)); vector<int>vpar(N,-1); function<void(int,int)> DFS =[&](int cur,int par){ vpar[cur]= par;for(int i =0; i <20; i++){ cnt[i][cur]=(i == c[cur]-1); preSum[i][cur]= cnt[i][cur];if(-1!= par){ preSum[i][cur]+= preSum[i][par];}}for(constauto& next : neiBo[cur]){if(next == par){continue;}DFS(next, cur);}};DFS(0,-1);auto leves =CBFSLeve::Leve(neiBo,{0}); C2Parents p2(vpar, leves); vector<int> ans;for(auto[u, v]: que){ u--, v--;constint g = p2.GetPublicParent(u, v);int cur =0;for(int i =0; i <20; i++){ cur +=(preSum[i][u]+ preSum[i][v]-2* preSum[i][g]+ cnt[i][g]>0);} ans.emplace_back(cur);}return ans;}};intmain(){#ifdef_DEBUGfreopen("a.in","r",stdin);#endif// DEBUG  ios::sync_with_stdio(0); cin.tie(nullptr);//CInBuff<> in; COutBuff<10'000'000> ob;int N, Q; cin >> N >> Q;auto c = Read<int>(N);auto edge = Read<pair<int,int>>(N -1);auto que = Read<pair<int,int>>(Q);#ifdef_DEBUG printf("N=%d", N);Out(c,",c=");Out(que,",que=");Out(edge,",edge=");//Out(edge2, ",edge2=");//Out(rr, ",rr=");//Out(ab, ",ab=");//Out(par, "par=");//Out(que, "que=");//Out(B, "B=");#endif// DEBUG  Solution slu;auto res = slu.Ans(N,c,edge,que);for(constauto& i : res){ cout << i <<"\n";}return0;};

单元测试

int N; vector<int> c; vector<pair<int,int>> edge, que;TEST_METHOD(TestMethod01){ N =4, c ={1,2,3,1}, que ={{4,3},{1,4}}, edge ={{1,2},{1,3},{2,4}};auto res =Solution().Ans(N, c, edge, que);AssertEx({3,2}, res);}

扩展阅读

视频课程

先学简单的课程，请移步ZEEKLOG学院，听白银讲师（也就是鄙人）的讲解。
https://edu.ZEEKLOG.net/course/detail/38771
如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.ZEEKLOG.net/lecturer/6176

测试环境

操作系统：win7 开发环境： VS2019 C++17
或者 操作系统：win10 开发环境： VS2022 C++17
如无特殊说明，本算法用**C++**实现。