PROBLEM LINK:
Setter: Hasan
Tester: Teja Vardhan Reddy
Editorialist: Taranpreet Singh
DIFFICULTY:
Easy-Medium
PREREQUISITES:
Observations, Maths (Chicken McNugget Theorem), Implementation.
PROBLEM:
Given an array A of size N, we need to rearrange the elements into an array B such that the arithmetic average is the same for every subarray of B of length K.
QUICK EXPLANATION
- In array B, for every position p \geq K, B[p] = B[p-K] holds.
- If N is divisible by K, there exists a solution if and only if the frequency of each element is divisible by N/K, since each element needs to appear exactly N/K times.
- Otherwise, the first N \% K elements in B array shall be repeated N/K+1 times, and K- N\%K elements shall be repeated N/K times.
- For every value with frequency F, we can find such non-negative x, y, z such that (N/K+1)*x+(N/K)*(N/K+1)*y+(N/K)*z = F while maximizing. Now, we can fill the B array by splitting y as required into the first and second group as required.
- If no such x, y, z exists, there is no valid distribution.
EXPLANATION
First of all, let us see how the final array B looks like. Since all subarrays we are considering have size K, having the same arithmetic average for each subarray is equivalent to the condition that all subarrays of size K shall have the same sum.
Now considering first subarray sum S = \displaystyle\sum_{i = 0}^{K-1} B_i and second subarray sum T = \displaystyle\sum_{i = 1}^{K}B_i. We require S = T. On expanding S and T, we can see that B_1, B_2 \ldots B_{K-1} are present on both sides. So above condition reduces to B_0 = B_K.
Generalizing this, we can show that for every position p \geq K, the necessary and sufficient condition is B_p = B_{p-K}, which should hold in the final array.
Consider a simple case first where N is divisible by K.
In this case, we can split the array into blocks of size K and for each block, the first element should be same for all blocks, the second element should be same for all blocks and so on. We can see, that final array of size N is basically a repetition of an array of size K exactly N/K times. We need each element to occur in multiples of N/K.
This is it. If an element appears say X times where X is a multiple of N/K, then this element shall appear X/(N/K) times in each block.
Now assuming N is not a multiple of K. This point onwards, N/K denotes integer division here.
We still know that elements repeat in a group of K, so we need to choose elements from given elements such that the B_p = B_{p-K} holds for each $p \geq K.
But now, N not being divisible by K means that the last block is not a complete repetition of the first block. Specifically, first N \% K elements repeat N/K+1 times and next K-N \% K elements repeat exactly N/K times.
Let’s call a occurrences of an element as an a-group.
So, the final condition is, that we need N \% K (N/K+1)-groups and K - N\%K (N/K)-groups using elements of the array A.
Each value is independent here, as each group contains multiple occurrences of a single element. Let’s consider each element individually.
Say the frequency of the current element is F. We want to find such x, y, z \geq 0 (N/K+1)*x + (N/K)*(N/K+1)*y + (N/K)*z = F such that y is maximized. Let M = N/K, the equation become (M+1)*x+M*(M+1)*y+M*z =F.
The idea here is, that for each M*(M+1)-group, we can choose to split it into M (M+1)-groups, or (M+1) M-groups as per our needs.
Now assume we have fixed some y such that S = F - M*(M+1)*y \geq 0. Now we need (M+1)*x + M*z = S.
We can actually brute-force over all possible values of y in decreasing order and check for each y whether there’s a solution. It is guaranteed there can be at most one solution for a chosen y, see this theorem.
After finding such x, y and z, we can see, that current element minimum appear x times among N \%K positions, and minimum z times in remaining K - N\%K positions in the first block.
After finding all such x and z for each value, if \sum x exceeds N \%K or \sum z exceeds K - N \% K, then no solution is possible, since we cannot have more than N\%K (M+1)-groups, or more than K-N\%K M-groups.
Otherwise, after assigning these (M+1)-groups and M-groups, we now have y (M+1)*(M) groups for each value, which we can split into either (M+1) M-groups or M (M+1)-groups. If there’s no valid split exist, the answer is again NO, otherwise finding an assignment is easy, and that gives us the final first block, which we can repeat.
In case you are facing any issues in this step, refer the commented solution for details.
On a note, instead of trying all values, there are better ways, but this one is most intuitive. Feel free to share.
TIME COMPLEXITY
The time complexity is hard to ascertain for this approach, but I’ll give it a try.
In worst case, all elements appear roughly K*K+K times so that Internal loop is bound by K*K times but Chicken McNugget theorem states that for two co-prime numbers a and b, Exactly one of (x, (a-1)*(b-1)-x+1) can also be made, so this loop is gonna break much faster. I don’t have a proof, but seemingly it takes approximately O(K) with some constant. This is repeated for all (N/K) elements. There are also ordered maps involved, so that’s another O(log(K)) involved (size of the map never exceeds K), resulting in complexity O(N*log(K)) with a constant.
Do comment, if you have a more rigorous proof, proving some lower bound.
SOLUTIONS:
Tester's Solution
//raja1999
//#pragma comment(linker, "/stack:200000000")
//#pragma GCC optimize("Ofast")
//#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2")
#include <bits/stdc++.h>
#include <vector>
#include <set>
#include <map>
#include <string>
#include <cstdio>
#include <cstdlib>
#include <climits>
#include <utility>
#include <algorithm>
#include <cmath>
#include <queue>
#include <stack>
#include <iomanip>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
//setbase - cout << setbase (16)a; cout << 100 << endl; Prints 64
//setfill - cout << setfill ('x') << setw (5); cout << 77 <<endl;prints xxx77
//setprecision - cout << setprecision (14) << f << endl; Prints x.xxxx
//cout.precision(x) cout<<fixed<<val; // prints x digits after decimal in val
using namespace std;
using namespace __gnu_pbds;
#define f(i,a,b) for(i=a;i<b;i++)
#define rep(i,n) f(i,0,n)
#define fd(i,a,b) for(i=a;i>=b;i--)
#define pb push_back
#define mp make_pair
#define vi vector< int >
#define vl vector< ll >
#define ss second
#define ff first
#define ll long long
#define pii pair< int,int >
#define pll pair< ll,ll >
#define sz(a) a.size()
#define inf (1000*1000*1000+5)
#define all(a) a.begin(),a.end()
#define tri pair<int,pii>
#define vii vector<pii>
#define vll vector<pll>
#define viii vector<tri>
#define mod (1000*1000*1000+7)
#define pqueue priority_queue< int >
#define pdqueue priority_queue< int,vi ,greater< int > >
#define int ll
typedef tree<
int,
null_type,
less<int>,
rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;
//std::ios::sync_with_stdio(false);
vi vec;
int a[100005],arr[100005];
vii cnt1;
vii :: iterator it1;
set<pii> cnt;
set<pii>::iterator it;
main(){
std::ios::sync_with_stdio(false); cin.tie(NULL);
int t,t1;
cin>>t;
t1=t;
while(t--){
int n,k,i,c,c1,h,j,id,val,val1,fl=0,val2,c2=0;
vec.clear();
cnt1.clear();
cnt.clear();
cin>>n>>k;
rep(i,n){
cin>>a[i];
vec.pb(a[i]);
}
h=n/k;
c=k;
if(n%k){
h++;
c=n%k;
}
c1=k-c;
id=n%k;
sort(all(vec));
c2=1;
f(i,1,vec.size()){
if(vec[i]!=vec[i-1]){
if(c2%h==0)
cnt.insert(mp(c2,vec[i-1]));
else
cnt1.pb(mp(c2,vec[i-1]));
c2=1;
}
else{
c2++;
}
}
if(c2%h==0)
cnt.insert(mp(c2,vec[i-1]));
else
cnt1.pb(mp(c2,vec[i-1]));
for(it1=cnt1.begin();it1!=cnt1.end();it1++){
while(c1>0&&it1->ff>=h-1&&it1->ff%h){
it1->ff-=h-1;
j=id;
while(j<n){
arr[j]=it1->ss;
j+=k;
}
c1--;
id++;
}
if(it1->ff%h){
fl=1;
break;
}
if(it1->ff!=0){
cnt.insert(*it1);
}
}
if(fl||c1%h){
cout<<"NO"<<endl;
continue;
}
while(c1>0){
it=cnt.end();
it--;
val=it->ff;
if(val<h*(h-1)){
fl=1;
break;
}
val2=it->ff;
val=it->ss;
val1=h;
while(val1--){
j=id;
while(j<n){
arr[j]=val;
j+=k;
}
id++;
}
c1-=h;
cnt.erase(it);
cnt.insert(mp(val2-(h*(h-1)),val));
}
if(fl){
cout<<"NO"<<endl;
continue;
}
id=0;
for(it=cnt.begin();it!=cnt.end();it++){
val=it->ff;
while(val>0){
val-=h;
j=id;
while(j<n){
arr[j]=it->ss;
j+=k;
}
id++;
}
}
cout<<"YES"<<endl;
rep(i,n){
cout<<arr[i]<<" ";
}
cout<<endl;
}
return 0;
}
Editorialist's Solution
import java.util.*;
import java.io.*;
import java.text.*;
class EQAVG{
//SOLUTION BEGIN
void pre() throws Exception{}
void solve(int TC) throws Exception{
int n = ni(), k = ni();
int[] a = new int[n];
TreeMap<Integer, Integer> mp = new TreeMap<>();
for(int i = 0; i< n; i++){
a[i] = ni();
mp.put(a[i], mp.getOrDefault(a[i], 0)+1);
}
boolean valid = true;
int[] ans = new int[k];//If an answer exist, first k elements matter, since ans[i+x*k] = ans[i]
if(n%k == 0){
for(int i:mp.values())valid &= i%(n/k) == 0;//All elements should appear exactly c*(n/k) times for some constant c
if(valid){
int p = 0;
for(Map.Entry<Integer, Integer> e:mp.entrySet()){
int value = e.getKey(), freq = e.getValue();
while(freq > 0){
ans[p++] = value;
freq -= n/k;
}
}
}
}else{
int p1 = 0, p2 = n%k;//p0 -> start position to add (k+1)-freq-groups, p2 -> start position to add k-freq-groups
TreeMap<Integer, Integer> reserve = new TreeMap<>();
for(Map.Entry<Integer, Integer> e:mp.entrySet()){
int value = e.getKey(), freq = e.getValue();
int[] x = decompose(freq, n/k);
if(x == null){
valid = false;
break;
}
while(x[0] > 0 && p1 < n%k){
ans[p1++] = value;
x[0]--;
}
while(x[2] > 0 && p2 < k){
ans[p2++] = value;
x[2]--;
}
valid &= x[0] == 0 && x[2] == 0;//Checking for leftovers
reserve.put(value, x[1]);//Using (k+1)*k groups to utilize later as required.
}
if(valid){
for(Map.Entry<Integer, Integer> e:reserve.entrySet()){
int value = e.getKey(), freq = e.getValue();
{
int rem = (n%k-p1);//Number of remaining (n/k+1)-freq-groups required
//Checking whether we can use all n/k (n/k+1)-freq-groups generated from splitting one (n/k)*(n/k+1)-freq-group
while(freq > 0 && rem >= n/k){
rem -= n/k;
freq--;//Splitting this one group into (k+1) k-freq-group
for(int i = 0; i< n/k; i++)ans[p1++] = value;
}
}
{
int rem = (k-p2);//Number of remaining k-groups required.
//Checking whether we can use all k+1 k-freq-groups generated from splitting one k*(k+1)-freq-group
while(freq > 0 && rem >= n/k+1){
rem -= n/k+1;
freq--;
for(int i= 0; i< n/k+1; i++)ans[p2++] = value;
}
}
hold(freq == 0);
}
}
}
if(valid){
pn("YES");
for(int i = 0; i< n; i++)p(ans[i%k]+" ");pn("");
}else pn("NO");
}
//Splitting freq into 3 parts, such that (a+1)*x[0]+a*(a+1)*x[1]+a*x[2] = freq, and x[1] is maximized
//If no such solution exists, null is returned
int[] decompose(int freq, long a){
int max = (int)(freq/(a*(a+1)));
int[] x = new int[3];
for(x[1] = max; x[1] >= 0; x[1]--){
int rem = freq-(int)(x[1]*(a*(a+1)));
//Can we find such x[0] and x[2] such that x[0]*(k+1) + x[2]*k = rem
for(x[0] = 0; x[0]*(a+1) <= rem; x[0]++){
if((rem-x[0]*(a+1)) % a == 0){
x[2] = (int)((rem-x[0]*(a+1))/a);
return x;
}
}
}
//Turns out we can't
return null;
}
//SOLUTION END
void hold(boolean b)throws Exception{if(!b)throw new Exception("Hold right there, Sparky!");}
DecimalFormat df = new DecimalFormat("0.00000000000");
static boolean multipleTC = true;
FastReader in;PrintWriter out;
void run() throws Exception{
in = new FastReader();
out = new PrintWriter(System.out);
//Solution Credits: Taranpreet Singh
int T = (multipleTC)?ni():1;
pre();for(int t = 1; t<= T; t++)solve(t);
out.flush();
out.close();
}
public static void main(String[] args) throws Exception{
new EQAVG().run();
}
int bit(long n){return (n==0)?0:(1+bit(n&(n-1)));}
void p(Object o){out.print(o);}
void pn(Object o){out.println(o);}
void pni(Object o){out.println(o);out.flush();}
String n()throws Exception{return in.next();}
String nln()throws Exception{return in.nextLine();}
int ni()throws Exception{return Integer.parseInt(in.next());}
long nl()throws Exception{return Long.parseLong(in.next());}
double nd()throws Exception{return Double.parseDouble(in.next());}
class FastReader{
BufferedReader br;
StringTokenizer st;
public FastReader(){
br = new BufferedReader(new InputStreamReader(System.in));
}
public FastReader(String s) throws Exception{
br = new BufferedReader(new FileReader(s));
}
String next() throws Exception{
while (st == null || !st.hasMoreElements()){
try{
st = new StringTokenizer(br.readLine());
}catch (IOException e){
throw new Exception(e.toString());
}
}
return st.nextToken();
}
String nextLine() throws Exception{
String str = "";
try{
str = br.readLine();
}catch (IOException e){
throw new Exception(e.toString());
}
return str;
}
}
}
Feel free to share your approach, if you want to. (even if its same 
) . Suggestions are welcomed as always had been. 