// correct/job-amd.java import java.util.ArrayList; import java.util.TreeSet; class job { static final int maxn = 200_000 + 100; static long ans = 0L; ArrayList<Integer>[] adj = new ArrayList[maxn]; int[] u = new int[maxn]; int[] d = new int[maxn]; int[] p = new int[maxn]; static class Job implements Comparable<Job> { long u, d; private Object o; Job(int U, int D){ u = U; d = D; } void merge(Job j) { ans += d * j.u; u += j.u; d += j.d; } public int compareTo(Job j){ // -1: less that, 0: equal, +1: greater than long a = u * j.d; long b = j.u * d; return Long.compare(a, b); } public boolean equals(Object o) { if(!(o instanceof Job)) return false; Job j = (Job)o; return this.compareTo(j) == 0; } } TreeSet<Job>[] se = new TreeSet[maxn]; int dfs(int v){ ArrayList<Integer> children = new ArrayList<>(); int res = v; int big = -1; // index in children for(int i = 0; i < adj[v].size(); ++ i) { int u = adj[v].get(i); children.add(dfs(u)); if (big == -1 || se[children.get(i)].size() > se[children.get(big)].size()) big = i; } Job j = new Job(u[v], d[v]); if(big == -1){ // leaf se[res] = new TreeSet<>(); se[res].add(j); return res; } res = children.get(big); for(int i = 0; i < children.size(); ++ i) if(i != big) { for(Job jj: se[children.get(i)]) { Job it = se[res].ceiling(jj); if(it != null && jj.compareTo(it) == 0) { jj.merge(it); se[res].remove(it); } se[res].add(jj); } se[children.get(i)].clear(); } while(!se[res].isEmpty() && j.compareTo(se[res].last()) <= 0) { Job last = se[res].last(); j.merge(last); se[res].remove(last); } se[res].add(j); return res; } public long scheduling_cost(int[] P, int[] U, int[] D) { int n = U.length; for(int i = 0; i < n; ++ i) adj[i] = new ArrayList<>(); for(int i = 0; i < n; ++ i) { u[i] = U[i]; d[i] = D[i]; p[i] = P[i]; ans += (long)u[i] * (long)d[i]; if (p[i] != -1) adj[p[i]].add(i); } int x = dfs(0); long d = 0; while(!se[x].isEmpty()) { Job j = se[x].last(); ans += d * j.u; d += j.d; se[x].remove(j); } return ans; } }

Note: job.java uses unchecked or unsafe operations.
Note: Recompile with -Xlint:unchecked for details.