2016-2017 ACM-ICPC Southwestern European Regional F Performance Review

問題

Attachments - 2016-2017 ACM-ICPC Southwestern European Regional Programming Contest (SWERC 2016) - Codeforces

根つき木があって、各頂点にはランクとコストが定められています。各頂点について、その頂点の子孫であり、かつ、その頂点のランク未満のランクをもつ頂点のコストの合計を計算してください。

解法

オイラーツアーで子孫となる範囲を計算しておきます。こうすることで、各頂点について、子孫を数列上の範囲として管理できます。このような処理をしておくと、以下のようなクエリを処理する問題になります。

• 場所 pos に数字 cost を入れる。
• 範囲 (start, end) に含まれる数字の合計を計算する。

これは Fenwick tree を使うことで処理できます。

解法

import java.io.PrintWriter;
import java.util.*;

public class F {
  ArrayList<Integer>[] tree;

  int[] start;
  int[] end;
  private void solve(Scanner in, PrintWriter out) {
    int N = in.nextInt();
    tree = new ArrayList[N];
    for (int i = 0; i < N; i++) tree[i] = new ArrayList<>();

    start = new int[N];
    end = new int[N];
    int[] costs = new int[N];
    int[] ranks = new int[N];
    int rootId = -1;
    for (int i = 0; i < N; i++) {
      int parent = in.nextInt() - 1;
      int rank = in.nextInt();
      int cost = in.nextInt();

      costs[i] = cost;
      ranks[i] = rank;

      if (parent >= 0) {
        tree[i].add(parent);
        tree[parent].add(i);
      } else {
        rootId = i;
      }
    }

    dfs(rootId, -1);

    int[][] nodes = new int[N][2];
    for (int i = 0; i < N; i++) {
      nodes[i][0] = ranks[i];
      nodes[i][1] = i;
    }
    Arrays.sort(nodes, Comparator.comparingInt(o -> o[0]));

    FenwickTree bit = new FenwickTree(2 * N);
    long[] ans = new long[N];

    int addHead = 0;
    int sumHead = 0;
    while (addHead < N) {
      int curRank = nodes[addHead][0];
      while (addHead < N && nodes[addHead][0] == curRank) {
        int i = nodes[addHead][1];
        ans[i] = bit.sum(start[i], end[i]) / 2;
        addHead++;
      }
      while (sumHead < N && nodes[sumHead][0] == curRank) {
        int i = nodes[sumHead][1];
        bit.add(start[i], costs[i]);
        bit.add(end[i], costs[i]);
        sumHead++;
      }
    }
    for (long a : ans) out.println(a);
  }

  int now = 0;
  void dfs(int v, int p) {
    start[v] = now;
    now++;
    for (int u : tree[v]) {
      if (u == p) continue;
      dfs(u, v);
    }
    end[v] = now;
    now++;
  }

  public class FenwickTree {
    int N;
    long[] data;

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

    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);
    }
  }

  public static void main(String[] args) {
    Scanner in = new Scanner(System.in);
    PrintWriter out = new PrintWriter(System.out);
    new F().solve(in, out);
    in.close();
    out.close();
  }
}