【牛客练习赛 92】B 题题解
题目链接
题目大意
给定一个长度为n nn的数组a aa和一个正整数k kk，要求将数组a aa划分为k kk个互不相交的集合，且每个集合的元素和都不为0 00。
请构造满足条件的一种划分方案，如若不行输出NO \text{NO}NO。
数据范围
- 1 ≤ k ≤ n ≤ 1 0 5 , 1 \leq k \leq n \leq 10^5,1≤k≤n≤105,
- − 1 0 9 ≤ a i ≤ 1 0 9 . -10^9 \leq a_i \leq 10^9.−109≤ai≤109.
Solution
首先可以把所有0 00都放到第1 11个集合里，剩下的数再想办法加入这k kk个集合中的一个。
若k = 1 k = 1k=1，只能是∑ i = 1 n a i ≠ 0 \sum\limits_{i = 1}^{n}a_i \neq 0i=1∑nai=0才能满足条件。
若k > 1 k > 1k>1，那么考虑将数组a aa排序（除去0 00之后），若∣ a ∣ < k |a| < k∣a∣<k则无解。否则把前k − 1 k - 1k−1个数按序加入前k − 1 k - 1k−1个集合，最后一个集合加入a aa的末尾元素（即最大数）。
- 这样可以保证，如果同时存在正数和负数，它们不会在同一个集合里。如果直接把前k kk个数按序加入前k kk个集合，可能存在前k kk个数都是负数的情况，这样正数就没有自己的集合了。
最后对于a k , a k + 1 , ⋯ , a ∣ a ∣ − 1 a_k, a_{k + 1}, \cdots, a_{|a| - 1}ak,ak+1,⋯,a∣a∣−1，如果是负数，就加入第1 11个集合，否则加入第k kk个集合。
时间复杂度O ( n log n ) O(n\log n)O(nlogn)
- 瓶颈在排序，其实有O ( n ) O(n)O(n)的写法，但是懒了。
C++ Code
#include<bits/stdc++.h>usingi64=longlong;intmain(){std::ios::sync_with_stdio(false);std::cin.tie(nullptr);intn,k;std::cin>>n>>k;std::vector<int>a(n);for(auto&x:a){std::cin>>x;}intzer=0;std::vector<int>b;b.reserve(n);for(intx:a){if(x!=0){b.push_back(x);}else{zer++;}}a=std::move(b);std::ranges::sort(a);if(a.size()<k){std::cout<<"NO\n";return0;}std::vector<std::vector<int>>ans(k);ans[0]=std::vector(zer,0);autoprint=[&](){std::cout<<"YES\n";for(constauto&v:ans){std::cout<<v.size()<<" ";for(inti=0;i<v.size();i++){std::cout<<v[i]<<" \n"[i==v.size()-1];}assert(std::reduce(v.begin(),v.end(),0LL)!=0);}};if(k==1){ans[0]=a;if(std::reduce(a.begin(),a.end(),0LL)!=0){print();}else{std::cout<<"NO\n";}return0;}for(inti=0;i<k-1;i++){ans[i].push_back({a[i]});}ans.back().push_back(a.back());for(inti=k-1;i<a.size()-1;i++){if(a[i]<0){ans[0].push_back(a[i]);}else{ans.back().push_back(a[i]);}}print();return0;}