本文涉及知识点
C++算法：前缀和、前缀乘积、前缀异或的原理、源码及测试用例 包括课程视频
C++DFS
树上倍增 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>
using namespace std;
template<class T1, class T2>
std::istream& operator >> (std::istream& in, pair<T1, T2>& pr) {
in >> pr.first >> pr.second;
return in;
}
template<class T1, class T2, class T3 >
std::istream& operator >> (std::istream& in, tuple<T1, T2, T3>& t) {
in >> get<0>(t) >> get<1>(t) >> get<2>(t);
return in;
}
template<class T1, class T2, class T3, class T4 >
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<class T1, class T2, class T3, class T4, class T5, class T6, class T7 >
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<class T = 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<class T = int>
vector<T> ReadNotNum() {
vector<T> ret;
T tmp;
while (cin >> tmp) {
ret.emplace_back(tmp);
if ('\n' == cin.get()) { break; }
}
return ret;
}
template<class T = 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>
class COutBuff
{
public:
COutBuff() {
m_p = puffer;
}
template<class T>
void write(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();
}
void writestr(const char* sz) {
strcpy(m_p, sz);
m_p += strlen(sz);
AuotToFile();
}
inline void write(char ch)
{
*m_p++ = ch;
AuotToFile();
}
inline void ToFile() {
fwrite(puffer, 1, m_p - puffer, stdout);
m_p = puffer;
}
~COutBuff() {
ToFile();
}
private:
inline void AuotToFile() {
if (m_p - puffer > N - 100) {
ToFile();
}
}
char puffer[N], * m_p;
};
template<int N = 1'000'000>
class CInBuff
{
public:
inline CInBuff() {}
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();
int x(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>>(long long& val) {
FileToBuf();
long long x(0); int f(0);
while (!isdigit(*S))
f |= (*S++ == '-');
while (isdigit(*S))
x = (x << 1) + (x << 3) + (*S++ ^ 48);
val = f ? -x : x; S++;//忽略空格换行
return *this;
}
template<class T1, class T2>
inline CInBuff& operator>>(pair<T1, T2>& val) {
*this >> val.first >> val.second;
return *this;
}
template<class T1, class T2, class T3>
inline CInBuff& operator>>(tuple<T1, T2, T3>& val) {
*this >> get<0>(val) >> get<1>(val) >> get<2>(val);
return *this;
}
template<class T1, class T2, class T3, class T4>
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<class T = 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<class T = int>
vector<T> Read(int n) {
vector<T> ret(n);
for (int i = 0; i < n; i++) {
*this >> ret[i];
}
return ret;
}
template<class T = int>
vector<T> Read() {
vector<T> ret;
*this >> ret;
return ret;
}
private:
inline void FileToBuf() {
const int 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;
};
class CNeiBo
{
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 (const auto& [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 (const auto& 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 (const auto& 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 (const auto& [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;
}
};
class CBFSLeve {
public:
static vector<int> Leve(const vector<vector<int>>& neiBo, vector<int> start) {
vector<int> leves(neiBo.size(), -1);
for (const auto& s : start) {
leves[s] = 0;
}
for (int i = 0; i < start.size(); i++) {
for (const auto& next : neiBo[start[i]]) {
if (-1 != leves[next]) { continue; }
leves[next] = leves[start[i]] + 1;
start.emplace_back(next);
}
}
return leves;
}
template<class NextFun>
static vector<int> Leve(int N, NextFun nextFun, vector<int> start) {
vector<int> leves(N, -1);
for (const auto& s : start) {
leves[s] = 0;
}
for (int i = 0; i < start.size(); i++) {
auto nexts = nextFun(start[i]);
for (const auto& 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) {
const int 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) {
const int 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 (const auto& v : leveNodes) {
ret.insert(ret.end(), v.begin(), v.end());
}
return ret;
};
};
class CParents
{
public:
CParents(vector<int>& vParent, long long iMaxDepth)
{
int iBitNum = 0;
for (; iMaxDepth; iBitNum++) {
const auto mask = 1LL << iBitNum;
if (mask & iMaxDepth) { iMaxDepth = iMaxDepth ^ mask; }
}
const int 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++)
{
const int iPre = m_vParents[i - 1][j];
if (-1 != iPre)
{
m_vParents[i][j] = m_vParents[i - 1][iPre];
}
}
}
}
int GetParent(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;
}
inline int GetBitCnt()const { return m_vParents.size(); };
inline const int& GetPow2Parent(int iNode, int n)const {
return m_vParents[n][iNode];
}
//在leftNodeExclude的1到rightLeve级祖先中查找符合pr的最近祖先
template<class _Pr>
int FindFirst(int leftNodeExclude, int rightLeve, _Pr pr) {
for (int iBit = GetBitCnt() - 1; iBit >= 0; iBit--) {
const int iMask = 1 << iBit;
if (!(iMask & rightLeve)) { continue; }
if (pr(m_vParents[iBit][leftNodeExclude])) {
return BFindFirst(leftNodeExclude, iBit, pr);
}
leftNodeExclude = m_vParents[iBit][leftNodeExclude];
}
return -1;
}
//在node的0到rightLeve级祖先中查找符合pr的最远祖先比node高多少层次,这些层次必须存在
template<class _Pr>
int FindEnd(int node, int rightLeve, _Pr pr) {
int leve = 0;
for (int iBit = GetBitCnt() - 1; iBit >= 0; iBit--) {
const int 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>
int BFindFirst(int leftNodeExclude, int pow, _Pr pr) {
while (pow > 0) {
const int& 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>
int BFindEnd(int node, int pow, _Pr pr) {
int leve = 0;
while (pow > 0) {
pow--;
const int& mid = m_vParents[pow][node];
if (pr(mid)) {
node = mid;
leve = leve ^ (1 << pow);
}
else {
}
}
return leve;
}
vector<vector<int>> m_vParents;
};
class C2Parents : public CParents
{
public:
C2Parents(vector<int>& vParent, const vector<int>& vDepth) :m_vDepth(vDepth)
, CParents(vParent, *std::max_element(vDepth.begin(), vDepth.end()))
{
}
int GetPublicParent(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--) {
const int iMask = 1 << iBit;
if (iMask & leve0) {
const int i1 = GetPow2Parent(iNode1, iBit);
const int i2 = GetPow2Parent(iNode2, iBit);
if (i1 == i2) {
while (iBit > 0) {
const int i3 = GetPow2Parent(iNode1, iBit - 1);
const int i4 = GetPow2Parent(iNode2, iBit - 1);
if (i3 != i4) {
iNode1 = i3; iNode2 = i4;
}
iBit--;
}
return GetPow2Parent(iNode1, 0);
}
else {
iNode1 = i1; iNode2 = i2; leve0 -= iMask;
}
}
}
return iNode1;
}
protected:
vector<vector<int>> m_vParents;
const vector<int> m_vDepth;
};
class Solution {
public:
vector<int> Ans(const int 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 (const auto& 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--;
const int 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;
}
};
int main() {
#ifdef _DEBUG
freopen("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 (const auto& i : res)
{
cout << i << "\n";
}
return 0;
};
单元测试
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);
}
扩展阅读
| 我想对大家说的话 |
|---|
| 工作中遇到的问题，可以按类别查阅鄙人的算法文章，请点击《算法与数据汇总》。 |
| 学习算法：按章节学习《喜缺全书算法册》，大量的题目和测试用例，打包下载。重视操作 |
| 有效学习：明确的目标 及时的反馈 拉伸区（难度合适） 专注 |
| 闻缺陷则喜(喜缺)是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。 |
| 子墨子言之：事无终始，无务多业。也就是我们常说的专业的人做专业的事。 |
| 如果程序是一条龙，那算法就是他的是睛 |
| 失败+反思=成功 成功+反思=成功 |
视频课程
先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771
如何你想快速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.csdn.net/lecturer/6176
测试环境
操作系统：win7 开发环境： VS2019 C++17
或者 操作系统：win10 开发环境： VS2022 C++17
如无特殊说明，本算法用**C++**实现。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
