Tech Vineyard: Codejam Chain Reactions

Introduction

This article describes a Java API to run a Depth-First-Search traversal. It lets you customize node processing via a node visitor.

Problem

In Code Jam Chain Reactions problem, we need to implement DFS to compute a root node value. It is derived from its children values in a recursive fashion.

Constraints

1 ≤ T ≤ 10^2
1 ≤ N ≤ 10^5

Input Parsing

Nodes are indexed 1 through N. We assign index 0 to the abyss, which acts as a sink.

package codejam.qualification2022;

import graph.tree.traversal.DFSTraversal;
import graph.tree.traversal.NodeVisitor;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;

public class ChainReactions {
   public static void main(String[] args) throws Exception {
      //InputStream inputStream = System.in;
      InputStream inputStream = new FileInputStream("ChainReactions");
      BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
      PrintWriter writer = new PrintWriter(new BufferedOutputStream(System.out));

      String[] tokens;
      tokens = bufferedReader.readLine().split(" ");
      int T = Integer.parseInt(tokens[0]);
      for (int t = 1; t <= T; t++) {
         tokens = bufferedReader.readLine().split(" ");
         int N = Integer.parseInt(tokens[0]);
         tokens = bufferedReader.readLine().split(" ");
         int[] F = new int[N+1];
         for (int i = 1; i <= N; i++) {
            F[i] = Integer.parseInt(tokens[i-1]);
         }
         tokens = bufferedReader.readLine().split(" ");
         int[] P = new int[N+1];
         for (int i = 1; i <= N; i++) {
            P[i] = Integer.parseInt(tokens[i-1]);
         }
         ChainReactions chainReactions = new ChainReactions();
         writer.println(String.format("Case #%d: %s", t, chainReactions.maxFun(N, F, P)));
      }
      writer.close();
      inputStream.close();
   }
}

Input

3
4
60 20 40 50
0 1 1 2
5
3 2 1 4 5
0 1 1 1 0
8
100 100 100 90 80 100 90 100
0 1 2 1 2 3 1 3

Output

Case #1: 110
Case #2: 14
Case #3: 490

Adjacency List

Before running graph traversal, we need to build adjacency list from parent P array.

private List<List<Integer>> getAdjacency(int N, int[] P) {
   List<List<Integer>> adjacency = new ArrayList<>();
   for (int i = 0; i <= N; i++) {
      adjacency.add(new ArrayList<>());
   }
   for (int i = 1; i <= N; i++) {
      adjacency.get(P[i]).add(i);
   }
   return adjacency;
}

It defines a tree structure, which is a directed acyclic connected graph. No cycles are present per the assumption that a node of index i can only have a parent of index j < i. Arrow direction is reverted in the traversal compared to the problem description: In the adjacency list, a node always points to its list of children.

From the problem description,

Each module may point at one other module with a lower index.

Solution

We add each cluster maxFun values..

Clusters are paths of nodes starting at any leaves, delimited optimally to maximize the sum of cluster fun attribute.

private static final int ROOT = 0;

public long maxFun(int N, int[] F, int[] P) {
   List<List<Integer>> adjacency = getAdjacency(N, P);

   FunNodeVisitor visitor = new FunNodeVisitor();
   visitor.F = F;
   visitor.maxFun = new int[N+1];

   visitor.fun = 0;
   new DFSTraversal().traverse(adjacency, visitor);
   visitor.fun += visitor.maxFun[ROOT];

   return visitor.fun;
}

Node visitor

We define following interface

package graph.tree.traversal;

import java.util.List;

public interface NodeVisitor {
   void visit(int parent, List<Integer> children);
}

In the solution's implementation, we sum up the maxFun values of each children, except the child with the minimum value. In this case we assign the max of current nodes's F value and the min child maxFun value.

static class FunNodeVisitor implements NodeVisitor {
   int[] F;
   int[] maxFun;
   long fun;

   public void visit(int parent, List<Integer> children) {
      // compute node value bottom-up
      int minChild = minFunChild(children, maxFun);
      maxFun[parent] = Math.max(minChild == -1 ? 0 : maxFun[minChild], F[parent]);

      fun += children.stream()
            .filter(child -> child != minChild)
            .mapToLong(child -> maxFun[child])
            .sum();
   }

   private int minFunChild(List<Integer> children, int[] maxFun) {
      return children.stream()
            .min(Comparator.comparingInt(child -> maxFun[child]))
            .orElse(-1);
   }
}

DFS traversal

We define a general traversal interface:

package graph.tree.traversal;

import java.util.List;

public interface Traversal {
   void traverse(List<List<Integer>> adjacency, NodeVisitor nodeVisitor);
}

Traversal will explore all the nodes in the adjacency list.

Algorithm

A DFS reference can be found in CLRS, chapter 22, section 3.

Exercise 22.3-6 asks for the Stack implementation,

Rewrite the procedure DFS, using a stack to eliminate recursion.

Implementation

See this commit for the full implementation details.

package graph.tree.traversal;

import java.util.Deque;
import java.util.LinkedList;
import java.util.List;

/**
 * Simulate a Depth-First-Search Postorder traversal
 * Use a LIFO queue where node elements are visited twice:
 * - 1st time to explore the children top-down
 * - 2nd time to provide a bottom-up node processing
 */
public class DFSTraversal implements Traversal {

   public void traverse(List<List<Integer>> adjacency, NodeVisitor nodeVisitor) {
      int N = adjacency.size();

      // Use a color attribute to visit each node twice during the queue processing
      // It also helps to discover connected components and detect loops
      Color[] color = new Color[N];
      for (int i = 0; i < N; i++) {
         color[i] = Color.WHITE;
      }

      for (int i = 0; i < N; i++) {
         if (color[i] == Color.WHITE) {
            // find the connected component
            traverse(adjacency, i, nodeVisitor, color);
         }
      }
   }

   private void traverse(List<List<Integer>> adjacency, int root, NodeVisitor nodeVisitor, Color[] color) {
      Deque<Integer> deque = new LinkedList<>();
      deque.addLast(root);

      while (! deque.isEmpty()) {
         // peek only on 1st visit, remove at the end of 2nd visit
         int parent = deque.getLast();
         List<Integer> c = adjacency.get(parent);
         if (color[parent] == Color.WHITE) {
            color[parent] = Color.GRAY;

            // explore tree top-down
            c.stream()
                  .filter(child -> color[child] == Color.WHITE) // avoid infinite loop
                  .forEach(deque::addLast);
         } else {
            // process node bottom up
            nodeVisitor.visit(parent, c);

            color[parent] = Color.BLACK;
            // poll on 2nd visit
            deque.removeLast();
         }
      }
   }
}

We use a Stack / LIFO queue to simulate DFS. This removes the tree depth limitation associated to recursion. Otherwise depending on your input you may run into stack overflow error. The LIFO order is generated via addLast / getLast / removeLast in Deque interface. We instantiate LinkedList class as our interface implementation choice.

The traversal notifies NodeVisitor callback on every node once. The are visited twice in the LIFO queue. We leverage the node color assignment to decide whether to explore or to backtrack.

Nodes are explored by adding them at the end of queue during the 1st top-down pass.
We can finally compute the parent value from its already computed children values, bottom-up, in the 2nd bottom-up pass. This logics is delegated to the NodeVisitor.

Unit Test

package graph.tree.traversal;

import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

import static java.util.Arrays.asList;

public class DFSTraversalTest {

   private final static Logger LOG = LoggerFactory.getLogger(DFSTraversalTest.class);

   private final DFSTraversal dfsTraversal = new DFSTraversal();

   @Test
   public void dfs() {
      List<List<Integer>> adjacency = new ArrayList<>();
      adjacency.add(0, asList(1));
      adjacency.add(1, asList(2, 3));
      adjacency.add(2, asList(4));
      adjacency.add(3, asList());
      adjacency.add(4, asList());

      List<Integer> nodes = new ArrayList<>();
      dfsTraversal.traverse(adjacency, (parent, children) -> {
         nodes.add(parent);
         LOG.info(String.format("%d: %s", parent, children));
      });
      Assertions.assertEquals(asList(3, 4, 2, 1, 0), nodes, "Post-order traversal");
   }
}

Here's the test logging output:

[]
[]
[4]
[2, 3]
[1]

Notice the last leaf is visited first, backtracking the post-order all the way back to the root.

Conclusion

Sometimes you just want to build on top of the construct without dealing with the implementation details. Once you know what type of traversal you need (DFS vs BFS), there are multiple benefits to get from those reusable components:

focus on custom node value computation logics to solve the problem at hand
use a well-tested, reusable and extensible implementation
reduce the risk of bugs associated to remembering the text book algorithm details
save time by not having to re-code the DFS implementation

Appendix: Graphviz input

Graph 1

digraph G {

subgraph cluster0 {

label="60"

labeljust=r

0 -> 1

1 -> 3

}

1 -> 2

subgraph cluster2 {

label="50"

labeljust=r

2 -> 4

}

0 [xlabel="0"]

1 [xlabel="60"]

2 [xlabel="20"]

3 [xlabel="40"]

4 [xlabel="50"]

}

Graph 2