paoPaoOptUtil.cc

#include <vector>
#include <unordered_map>
#include <map>
#include <cassert>
#include <iostream>

using namespace std;

class Solver {
public:
  Solver() = default;
  unordered_map<int, vector<int>> vToCnt(vector<vector<int>>& vals) {
    // get val -> ids
    // T: O(M), M size of vals
    unordered_map<int, vector<int>> res;
    for (auto val: vals) res[val[1]].push_back(val[0]);
    return res;
  }
  map<int, vector<int>> vToCntSorted(vector<vector<int>>& vals) {
    // get val -> ids, val sorted
    // T: O(MlgM), M size of vals
    map<int, vector<int>> res;
    for (auto val: vals) res[val[1]].push_back(val[0]);
    return res;
  }
  void merge(vector<int>& keys1, vector<int>& keys2, vector<vector<int>>& res) {
    for (int k1: keys1) {
      for (int k2: keys2) {
        res.push_back({k1, k2});
      }
    }
  }
  vector<vector<int>> solve(vector<vector<int>>& a, vector<vector<int>>& b, int target) {
    // T: O((N +M)lgM + M * N), N size of a, M size of b
    unordered_map<int, vector<int>> vaToCnt = vToCnt(a);
    map<int, vector<int>> vbToCnt = vToCntSorted(b);

    unordered_map<int, vector<vector<int>>> candidates;
    int closest;
    for (auto it = vaToCnt.begin(); it != vaToCnt.end(); it++) {
      int res = target - it->first;
      auto it2 = vbToCnt.upper_bound(res);
      if (it2 != vbToCnt.begin()) {
        it2--;
        int keySum = it->first + it2->first;
        if (candidates.empty()) closest = keySum;

        merge(it->second, it2->second, candidates[keySum]);
        keySum = max(keySum, closest);
      }
    }
    //cout << closest << endl;
    return candidates.empty() ? vector<vector<int>>() : candidates[closest];
  }
};

void tests() {
  {
    Solver s;
    vector<vector<int>> a = {{1, 2}, {2, 4}, {3, 6}};
    vector<vector<int>> b = {{1, 2}};
    int target = 7;
    vector<vector<int>> expected = {{ 2, 1 }};
    auto vals = s.solve(a, b, target);
    //for (auto v: vals) cout << v[0] << " " << v[1] << endl;
    assert(expected == s.solve(a, b, target));
  }
}

int main() {
  tests();
  return 1;
}