Polynomial Round 2022 (Div. 1 + Div. 2, Rated, Prizes!) D. Same Count One

Polynomial Round 2022 (Div. 1 + Div. 2, Rated, Prizes!) D. Same Count One

Considering that we need to make the number of $1$s in each array the same, we should calculate the sum of $1$s, and every array has $sum/n$ $1$s. Because only the same position of two different arrays can be selected for exchange each time, for a position pos , we traverse each array each time. If the number of $1$s in this array is not enough, then we need to turn some $0$s into $1$s; If the number of $1$s in this array is more than we need, then some $1$s should be turned into $0$s. It can be proved that as long as the total number of $1$s is a multiple of n, the number of $1$s in each array can be made the same through exchange.

#include<bits/stdc++.h>
using namespace std;

int main()
{
    int T;cin>>T;
    while(T--)
    {
        int n,m;cin>>n>>m;
        vector<vector<int>> a(n+1,vector<int>(m+1));
        multimap<int,int> b;
        int sum=0;
        for(int i=1;i<=n;i++)
        {
            int res=0;
            for(int j=1;j<=m;j++)
            {
                cin>>a[i][j];
                res+=a[i][j];
                sum+=a[i][j];
            }
            b.insert({res,i});
        }
        if(sum%n!=0)
        {
            cout<<"-1"<<endl;
            continue;
        }
        vector<array<int,3>> cnt;
        while(1)
        {
            auto begin=b.begin();
            auto end=--b.end();

            if(begin->first==end->first) break;
            else
            {
                int bx=begin->first,by=begin->second;
                int ex=end->first,ey=end->second;

                for(int i=1;i<=m;i++)
                    if(!a[by][i] && a[ey][i])
                    {
                        swap(a[by][i],a[ey][i]);
                        cnt.push_back({by,ey,i});
                        break;
                    }
                b.erase(begin);
                b.erase(end);
                b.insert({bx+1,by});
                b.insert({ex-1,ey});
            }
        }
        cout<<cnt.size()<<endl;
        for(auto &[x,y,z]:cnt)
            cout<<x<<" "<<y<<" "<<z<<endl;
    }
    return 0;
}