Analysis: Yin and Yang by Richard Peng

If we view Charcolais and Angus along an edge as weights of -1 or 1, this problem asks for the number of paths (with one or more edges) in a tree such that:

1. The total weight of the path is 0
2. There is a vertex in the middle of the path such that the weight to both ends is 0.

Note that the tree is unrooted, and can therefore be rooted arbitrarily. This allows to first consider the 'simpler' problem of counting the number of such paths that go through the root vertex in a highly efficient manner.

The case where the path ends at the root can be checked in O(n) time. Otherwise, the path is uniquely defined by its two end points, which must use edges leaving the root along different edges. Therefore, if we only want the first condition, we can loop over the first end point. If the weight of this path to root is w, then the weight of the path from the root to the second end point must be -w. A look up array then lets us count the number of such second end points in O(1) time. One detail omitted here is that an extra look up chart (or a more careful sweeping order through the children) is needed to deal with the over-counting caused by the second endpoint using the same edge leaving the root.

Now consider the second condition of having a third point in the middle of the path such that the weight of the path to one of the endpoints is 0. Conceptually it might be easier to count the complement: paths where no intermediate point has a partial sum of 0. Since such a point must lie on the path between the root and one of its endpoints, it suffices to consider only paths from the root where the final distance (from the end point) does not occur in its prefix. This condition can be incorporated in the DFS used to generate such paths using the same lookup table.

Now we've shown how to handle paths through a root in O(n) time, recall that the tree can be rooted arbitrarily. This fact is crucial since every tree can be rooted in a way such that the subtrees of the root has size at most 1/2 of the original size. Furthermore, such a root, known as a vertex separator, can be found in linear time using simple dynamic programming. Therefore if we pick the root as a separator, we can remove it and recurse on each of the resulting pieces. As the size of each piece halves per level of recursion, this leads to an O(nlogn) time algorithm. The below C++ solution demonstrates these ideas.

#include <iostream>
#include <vector>
#include <set>
#include <cstdio>

using namespace std;

typedef set<pair<int, int> >::const_iterator iter;

#define MAXN 100000

int tree_size[MAXN];
set<pair<int, int> > E[MAXN];

int compute_subtree_size(int u, int p) {
  int res = 1;
  for(iter i = E[u].begin(); i != E[u].end(); ++i) {
    if(i->first != p) {
      res += compute_subtree_size(i->first, u);
    }
  }
  return tree_size[u] = res;
}

int b_offset;
int c_offset;
pair<int, int> B[MAXN * 2 + 1];
pair<int, int> C[MAXN * 2 + 1];

void count_partials(int u, int p, int w, int mnw, int mxw) {
  if(w == mnw || w == mxw) {
    mnw = min(mnw, w - 1);
    mxw = max(mxw, w + 1);
    C[c_offset + w].second++;
  } else {
    C[c_offset + w].first++;
  }
  for(iter i = E[u].begin(); i != E[u].end(); ++i) {
    if(i->first != p) {
      count_partials(i->first, u, w + i->second, mnw, mxw);
    }
  }
}

/* Computes the number of balanced paths in the connected tree containing u. */
long long solve(int u) {
  /* Precompute the subtree size of all rooted subtrees of the tree. */
  int N = compute_subtree_size(u, -1);

  /* Move u down the tree until we find a node that is a separator, having no
   * child subtrees with more than half of the nodes. */
  int p = -1;
  for(;;) {
    bool separator = true;
    for(iter i = E[u].begin(); i != E[u].end(); ++i) {
      int v = i->first;
      if(v != p && tree_size[v] >= N / 2) {
        separator = false;
        p = u;
        u = v;
        break;
      }
    }
    if(separator) break;
  }

  /* We'll use u as the root of our tree.  First compute the number of balanced
   * paths going through u. */
  long long result = 0;

  /* Initialize our count data structure.  We offset the array as nodes with
   * negative partial sums will have a negative index.  B[offset + x].first is
   * the number of nodes processed so far with partial sum x and an intermediate
   * node on the way to root with the same partial sum.  .second is the same but
   * with no intermediate node. */
  b_offset = N;
  fill(B, B + 2 * N + 1, make_pair(0, 0));
  
  /* The root has partial sum 0 with no intermediate. */
  B[b_offset + 0].second = 1;

  for(iter i = E[u].begin(); i != E[u].end(); ++i) {
    int v = i->first;
    c_offset = tree_size[v];
    fill(C, C + 2 * tree_size[v] + 1, make_pair(0, 0));
    count_partials(v, u, i->second, i->second, i->second);

    for(int w = -tree_size[v]; w <= tree_size[v]; w++) {
      result += (long long)C[c_offset + w].first * B[b_offset - w].first;
      result += (long long)C[c_offset + w].first * B[b_offset - w].second;
      result += (long long)C[c_offset + w].second * B[b_offset - w].first;
    }

    /* Nodes with partial sums of 0 can use the root as an intermediate. */
    result += (long long)C[c_offset + 0].second * (B[b_offset + 0].second - 1);

    for(int w = -tree_size[v]; w <= tree_size[v]; w++) {
      B[b_offset + w].first += C[c_offset + w].first;
      B[b_offset + w].second += C[c_offset + w].second;
    }
  }
  
  /* Next, cut the edges coming from u and recursively solve the subtrees. */
  for(iter i = E[u].begin(); i != E[u].end(); ++i) {
    int v = i->first;
    E[v].erase(make_pair(u, i->second));
    result += solve(v);
  }

  return result;
}

int main() {
  freopen("yinyang.in", "r", stdin);
  freopen("yinyang.out", "w", stdout);

  /* Input tree. */
  int N; cin >> N;
  for(int i = 1; i < N; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    u--; v--;
    w = w == 0 ? -1 : 1; /* Adjust weight so it's -1/1 instead of 0/1. */

    E[u].insert(make_pair(v, w));
    E[v].insert(make_pair(u, w));
  }

  cout << solve(0) << endl;
  return 0;
}

There is also an alternate solution using the fact that the tables can be reused when we move the 'root' up the tree. There it's more convenient to only consider paths that go downwards from a vertex. Note that these weights can be tracked implicitly by taking the length from the actual root, and subtracting the length from the current 'root' for our calculations. Then all we need to do is to merge sets of weight values, which can be done in a total of O(nlog^2n) time if we always insert from the smaller set into the larger.