CS Academy Round #36 BBox Count

復習

問題

CS Academy

コード

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Scanner;

public class Main {

  private static final int MAX = 2500;

  private static long count(ArrayList<Integer> list, FenwickTree bit) {
    long ans = 0;
    int top;
    int bottom = 0;
    for (int p : list) {
      top = p;
      if (top - bottom >= 2) {
        ans += choose2Mod(bit.sum(bottom, top));
      }
      bottom = top + 1;
    }
    ans += choose2Mod(bit.sum(bottom, MAX));
    return ans;
  }

  private static long count(ArrayList<Integer> list1, ArrayList<Integer> list2, FenwickTree bit) {
    long ans = 0;
    int top;
    int bottom = 0;

    for (int i = 0, j = 0; i < list1.size() || j < list2.size(); ) {
      int p1 = i < list1.size() ? list1.get(i) : MAX;
      int p2 = j < list2.size() ? list2.get(j) : MAX;
      int p;
      if (p1 == p2) {
        p = p1;
        i++;
        j++;
      } else if (p1 > p2) {
        p = p2;
        j++;
      } else {
        p = p1;
        i++;
      }

      top = p;
      if (top - bottom >= 2) {
        ans += choose2Mod(bit.sum(bottom, top));
      }
      bottom = top + 1;
    }

    ans += choose2Mod(bit.sum(bottom, MAX));
    return ans;
  }

  public static void main(String[] args) {
    Scanner in = new Scanner(System.in);
    int N = in.nextInt();
    ArrayList<Integer>[] points = new ArrayList[MAX];
    for (int i = 0; i < MAX; i++) {
      points[i] = new ArrayList<>();
    }
    for (int i = 0; i < N; i++) {
      int x = in.nextInt() - 1;
      int y = in.nextInt() - 1;
      points[x].add(y);
    }
    for (ArrayList<Integer> p : points) {
      Collections.sort(p);
    }

    long ans = 0;

    boolean[] added = new boolean[MAX];
    FenwickTree bit = new FenwickTree(MAX);
    for (int i = 0; i < MAX; i++) {
      if (points[i].isEmpty()) {
        continue;
      }

      Arrays.fill(added, false);
      bit.clear();

      for (int j = i; j < MAX; j++) {
        if (points[j].isEmpty()) {
          continue;
        }
        for (int p : points[j]) {
          if (!added[p]) {
            bit.add(p, 1);
            added[p] = true;
          }
        }

        if (i == j) {
          continue;
        }
        ArrayList<Integer> left = points[i];
        ArrayList<Integer> right = points[j];

        long total = choose2Mod(bit.sum(0, MAX));
        ans += total;
        ans -= count(left, bit);
        ans -= count(right, bit);
        ans += count(left, right, bit);
      }
    }

    System.out.println(ans);
  }


  private static long choose2Mod(long x) {
    return (x * (x - 1) >> 1);
  }

  static class FenwickTree {

    int N;
    long[] data;

    FenwickTree(int N) {
      this.N = N + 1;
      data = new long[N + 1];
    }

    void clear() {
      Arrays.fill(data, 0);
    }

    void add(int k, long val) {
      for (int x = k; x < N; x |= x + 1) {
        data[x] += val;
      }
    }

    // [0, k)
    long sum(int k) {
      if (k >= N) {
        k = N - 1;
      }
      long ret = 0;
      for (int x = k - 1; x >= 0; x = (x & (x + 1)) - 1) {
        ret += data[x];
      }
      return ret;
    }

    // [l, r)
    long sum(int l, int r) {
      return sum(r) - sum(l);
    }
  }
}