1886E - I Wanna be the Team Leader - CodeForces Solution

#include <iostream>
#include <map>
#include <unordered_map>
#include <vector>
#include <unordered_set>
#include <set>
#include <queue>
#include <utility>
#include <algorithm>
#include <cassert>

using namespace std;
typedef long long ll;
typedef pair<ll, ll> pi;
typedef vector<ll> vi;
typedef vector<pi> vpi;


const int S = 2e5 + 7, M = 20;
const int mod = 1e9 + 7;
int b[S], n, m, suf[M][S], dp[1 << M];
vi tip[1 << M], ans[M];
vpi a;
// inject here

void calc_suf(int p) {
    for (int i = 0; i <= n; i ++){
        suf[p][i] = -1;
    }
    for (int en = n, beg = n; en; en --) {
        while (beg > en) beg --;
        while (b[p]  > (ll) a[beg].first * (en - beg + 1)) {
            if (beg-- == 1) return;
        }
        suf[p][en] = en - beg + 1;
    }
}

void solve(){
    cin >> n >> m;
    a.push_back({0, -1});
    for (int i = 1; i <= n; i ++){
        int x;
        cin >> x;
        a.push_back({x, i});
    }
    sort(a.begin(), a.end());

    for (int j = 0; j < m; j ++){
        cin >> b[j];
    }
    for (int i = 0; i < m; i ++){
        calc_suf(i);
        // cout << i << ": ";
        // for (int j = 1; j <= n; j ++) {
        //     cout << suf[i][j] << " ";
        // }
        // cout << endl;
        
    }
        
    for (int i = 1; i < (1 << m); i ++){
        dp[i] = mod;
        for (int j = 0; j < m; j ++) {
            if (!(i & (1 << j))) continue;
            int sub_mask = i - (1 << j);
            if (dp[sub_mask] == mod) continue;
            int jump = suf[j][n - dp[sub_mask]];
            if (jump == -1) continue;
            int cand = dp[sub_mask] + jump;
            if (dp[i] > cand) {
                dp[i] = cand;
                         // bit, start, end, next_mask
                tip[i] = {j, n - cand + 1, n - cand + jump, sub_mask};
            }   
        }
        // cout << i << ": " << dp[i] << endl;
    }
    int cur_mask = (1 << m) - 1;
    if (dp[cur_mask] == mod) {
        cout << "NO\n";
        return;
    }
    cout << "YES\n";
    while (cur_mask) {
        // cout << "!" << " ";
        // for (auto e : tip[cur_mask]) cout << e << " ";
        // cout << endl;

        for (int i = tip[cur_mask][1]; i <= tip[cur_mask][2]; i ++){
            ans[tip[cur_mask][0]].push_back(a[i].second);
        }
        cur_mask = tip[cur_mask][3];    
    }
    for (int i = 0; i < m; i ++){
        cout << ans[i].size() << " ";
        for (auto e : ans[i]) cout << e << " ";
        cout << endl;
    }
}


int main(){
    ios_base::sync_with_stdio(0);
    int t = 1;
    // cin >> t;
    while (t --)
        solve();

    return 0;
}