Submission #956787

#	Time	Username	Problem	Language	Result	Execution time	Memory
956787		chrisvilches	Dragon 2 (JOI17_dragon2)	C++14	100 / 100	695 ms	29868 KiB

dragon2

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

struct Point {
  int x, y, idx, tribe;
  Point operator-(const Point p) const { return {x - p.x, y - p.y, idx, tribe}; }
  Point operator+(const Point p) const { return {x + p.x, y + p.y, idx, tribe}; }
  long long cross(const Point p) const { return x * (long long)p.y - y * (long long)p.x; }
  Point negate() const { return {-x, -y, idx, tribe}; }
};

template <class T, unsigned int N>
struct BIT {
  void clear(const int n) { memset(bit, 0, sizeof(T) * n); }

  T sum(int l, int r) { return sum(r) - sum(l - 1); }

  void add(unsigned int idx, const T delta) {
    for (; idx < N; idx = idx | (idx + 1)) bit[idx] += delta;
  }

 private:
  T bit[N];

  T sum(int r) {
    T ret = 0;
    for (; r >= 0; r = (r & (r + 1)) - 1) ret += bit[r];
    return ret;
  }
};

BIT<int, 30'001> bit;
Point A, B;
vector<unordered_map<int, int>> ans;
vector<vector<Point>> tribe_points, order_by_b;

bool above(const Point p) { return B.cross(p) > 0; }

Point to_upper(const Point p) { return above(p) ? p : p.negate(); }

bool operator<(const Point p, const Point q) {
  return to_upper(p).cross(to_upper(q)) > 0;
}

bool cmp_by_b(const Point p, const Point q) {
  const bool a1 = above(p);
  const bool a2 = above(q);
  if (a1 != a2) return a1;
  return p.cross(q) > 0;
}

// TODO: Make sure the namings make sense. How queries work is a bit confusing.
void query_attacked(const vector<Point>& attacked, const int attacker_idx) {
  const vector<Point>& attackers = tribe_points[attacker_idx];
  const vector<Point>& ord = order_by_b[attacker_idx];

  bit.clear(attackers.size());

  int j = 0;

  for (const Point p : attacked) {
    int total = 0;

    while (j < (int)attackers.size() && attackers[j] < p) {
      bit.add(attackers[j].idx, 1);
      j++;
    }

    const auto mid = lower_bound(ord.begin(), ord.end(), B.negate(), cmp_by_b);

    // TODO: Maybe prepopulating the lower half can simplify this code?
    if (above(p)) {
      const int a = lower_bound(ord.begin(), mid, p - B, cmp_by_b) - ord.begin();
      const auto it = lower_bound(mid, ord.end(), B - p, cmp_by_b);

      total += a - bit.sum(0, a - 1);
      total += it - mid;
      total -= bit.sum(mid - ord.begin(), it - ord.begin() - 1);

    } else {
      const int a = lower_bound(ord.begin(), ord.end(), p - B, cmp_by_b) - ord.begin();
      const int b = lower_bound(ord.begin(), mid, B - p, cmp_by_b) - ord.begin();

      total += bit.sum(a, ord.size() - 1);
      total += bit.sum(b, mid - ord.begin() - 1);
    }

    ans[attacker_idx][p.tribe] += total;
  }
}

void query_attacks(const vector<Point>& attackers, const int attacked_idx) {
  const vector<Point>& attacked = tribe_points[attacked_idx];
  const vector<Point>& ord = order_by_b[attacked_idx];

  bit.clear(attacked.size());

  for (const auto q : attacked) {
    if (!above(q)) bit.add(q.idx, 1);
  }

  int j = 0;

  for (const Point p : attackers) {
    while (j < (int)attacked.size() && attacked[j] < p) {
      const Point q = attacked[j];
      bit.add(q.idx, above(q) ? 1 : -1);
      j++;
    }

    const Point v = above(p) ? p - B : B - p;

    const int a = lower_bound(ord.begin(), ord.end(), v, cmp_by_b) - ord.begin();
    const int b = lower_bound(ord.begin(), ord.end(), v.negate(), cmp_by_b) - ord.begin();

    ans[p.tribe][attacked_idx] += bit.sum(a, b - 1);
  }
}

void handle_all_queries(const map<int, set<int>>& attacks,
                        const map<int, set<int>>& attacked) {
  vector<tuple<long long, int, bool>> queries;

  const auto add_weighted_query = [&](const auto& q, const bool is_attack) {
    const auto [tribe_idx, other] = q;
    const long long weight = tribe_points[tribe_idx].size() * other.size();
    queries.emplace_back(weight, tribe_idx, is_attack);
  };

  for (const auto& q : attacks) add_weighted_query(q, true);
  for (const auto& q : attacked) add_weighted_query(q, false);

  sort(queries.rbegin(), queries.rend());

  for (const auto& [_, tribe_idx, is_attack] : queries) {
    vector<Point> all_points;

    const auto& other_tribes = (is_attack ? attacks : attacked).at(tribe_idx);

    for (const int other_tribe_idx : other_tribes) {
      int i = tribe_idx, j = other_tribe_idx;
      if (!is_attack) swap(i, j);

      if (ans[i].count(j)) continue;

      for (const Point p : tribe_points[other_tribe_idx]) {
        all_points.emplace_back(p);
      }

      ans[i][j] = 0;
    }

    sort(all_points.begin(), all_points.end());

    // TODO: This is weird. Shouldn't it be is_attack -> query_attacks ????
    if (is_attack) {
      query_attacked(all_points, tribe_idx);
    } else {
      query_attacks(all_points, tribe_idx);
    }
  }
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  int N, M, Q;

  while (cin >> N >> M) {
    tribe_points.assign(M + 1, vector<Point>());
    order_by_b.assign(M + 1, vector<Point>());
    ans.assign(M + 1, unordered_map<int, int>());

    for (int i = 0; i < N; i++) {
      Point p;
      cin >> p.x >> p.y >> p.tribe;
      tribe_points[p.tribe].push_back(p);
    }

    cin >> A.x >> A.y >> B.x >> B.y;
    cin >> Q;
    B = B - A;

    for (int m = 0; m <= M; m++) {
      auto& points = tribe_points[m];
      for (auto& p : points) p = p - A;
      for (auto& p : points) p = p - B;
      sort(points.begin(), points.end(), cmp_by_b);
      for (int i = 0; i < (int)points.size(); i++) points[i].idx = i;
      auto& b = order_by_b[m];
      b = points;
      for (auto& p : points) p = p + B;
      sort(points.begin(), points.end());
    }

    map<int, set<int>> attacks, attacked;
    vector<pair<int, int>> query_original_order;

    while (Q--) {
      int i, j;
      cin >> i >> j;
      attacks[i].emplace(j);
      attacked[j].emplace(i);
      query_original_order.emplace_back(i, j);
    }

    handle_all_queries(attacks, attacked);

    for (const auto& [i, j] : query_original_order) {
      cout << ans[i][j] << '\n';
    }
  }
}

Compilation message (stderr)

dragon2.cpp: In lambda function:
dragon2.cpp:125:16: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  125 |     const auto [tribe_idx, other] = q;
      |                ^
dragon2.cpp: In function 'void handle_all_queries(const std::map<int, std::set<int> >&, const std::map<int, std::set<int> >&)':
dragon2.cpp:135:20: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  135 |   for (const auto& [_, tribe_idx, is_attack] : queries) {
      |                    ^
dragon2.cpp: In function 'int main()':
dragon2.cpp:209:22: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  209 |     for (const auto& [i, j] : query_original_order) {
      |                      ^

• Subtask 1: 15 / 15
• Subtask 2: 45 / 45
• Subtask 3: 40 / 40