Java JTree expand by depth

import java.awt.Rectangle;

import javax.swing.JFrame;
import javax.swing.JScrollPane;
import javax.swing.JTree;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.TreePath;
import javax.swing.tree.TreeSelectionModel;

public class Main {
   public static void main(String[] args) {
      DefaultMutableTreeNode rootNode = new DefaultMutableTreeNode("Root");
      DefaultMutableTreeNode a = new DefaultMutableTreeNode("A");
      rootNode.add(a);//from   ww w . j a  va2 s .  c  om

      DefaultMutableTreeNode a1 = new DefaultMutableTreeNode("1");
      a.add(a1);
      DefaultMutableTreeNode a2 = new DefaultMutableTreeNode("2");
      a.add(a2);

      DefaultMutableTreeNode b = new DefaultMutableTreeNode("B");
      rootNode.add(b);
      DefaultMutableTreeNode b1 = new DefaultMutableTreeNode("1");
      b.add(b1);

      JTree tree = new JTree(rootNode);

      tree.getSelectionModel().setSelectionMode(TreeSelectionModel.SINGLE_TREE_SELECTION);

      JFrame f = new JFrame();
      f.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
      f.add(new JScrollPane(tree));
      f.setSize(250, 250);
      f.setVisible(true);

      expandJTree(tree, 1);
   }

   /**
    * Expands all nodes in a JTree.
    *
    * @param tree  The JTree to expand.
    * @param depth The depth to which the tree should be expanded.  Zero
    *              will just expand the root node, a negative value will
    *              fully expand the tree, and a positive value will
    *              recursively expand the tree to that depth.
    */
   public static void expandJTree(javax.swing.JTree tree, int depth) {
       javax.swing.tree.TreeModel model = tree.getModel();
       expandJTreeNode(tree, model, model.getRoot(), 0, depth);
   }

   /**
    * Expands a given node in a JTree.
    *
    * @param tree  The JTree to expand.
    * @param model The TreeModel for tree.
    * @param node  The node within tree to expand.
    * @param row   The displayed row in tree that represents
    *              node.
    * @param depth The depth to which the tree should be expanded.
    *              Zero will just expand node, a negative
    *              value will fully expand the tree, and a positive
    *              value will recursively expand the tree to that
    *              depth relative to node.
    */
   public static int expandJTreeNode(javax.swing.JTree tree,
           javax.swing.tree.TreeModel model, Object node, int row,
           int depth) {
       if (node != null && !model.isLeaf(node)) {
           tree.expandRow(row);
           if (depth != 0) {
               for (int index = 0; row + 1 < tree.getRowCount()
                       && index < model.getChildCount(node); index++) {
                   row++;
                   Object child = model.getChild(node, index);
                   if (child == null)
                       break;
                   javax.swing.tree.TreePath path;
                   while ((path = tree.getPathForRow(row)) != null
                           && path.getLastPathComponent() != child)
                       row++;
                   if (path == null)
                       break;
                   row = expandJTreeNode(tree, model, child, row,
                           depth - 1);
               }
           }
       }
       return row;
   }
}