Java JTree collapse all node

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);// w  w  w  .  jav  a  2 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);

      int s = getVisibleRowCount(tree);
      System.out.println(s);
   }

   public static void collapseAll(JTree tree, int keepSelectionLevel) {
      TreePath leadSelectionPath = tree.getLeadSelectionPath();
      // Collapse all
      int row = tree.getRowCount() - 1;
      while (row >= 0) {
         tree.collapseRow(row);
         row--;
      }
      DefaultMutableTreeNode root = (DefaultMutableTreeNode) tree.getModel().getRoot();
      tree.expandPath(new TreePath(root));
      if (leadSelectionPath != null) {
         Object[] path = leadSelectionPath.getPath();
         Object[] pathToSelect = new Object[path.length > keepSelectionLevel ? keepSelectionLevel : path.length];
         System.arraycopy(path, 0, pathToSelect, 0, pathToSelect.length);
         selectPath(tree, new TreePath(pathToSelect));
      }
   }

   public static void selectPath(JTree tree, TreePath path) {
      tree.makeVisible(path);
      showRowCentred(tree, tree.getRowForPath(path));
   }

   private static void showRowCentred(JTree tree, int row) {
      showRowCentred(tree, row, true);
   }

   public static void showRowCentred(JTree tree, int row, boolean centerHorizontally) {
      int visible = getVisibleRowCount(tree);
      int top = visible > 0 ? row - (visible - 1) / 2 : row;
      int bottom = visible > 0 ? top + visible - 1 : row;
      showAndSelect(tree, top, bottom, row, centerHorizontally);
   }

   private static int getVisibleRowCount(JTree tree) {
      Rectangle visible = tree.getVisibleRect();
      int count = 0;
      for (int i = 0; i < tree.getRowCount(); i++) {
         Rectangle bounds = tree.getRowBounds(i);
         if (visible.y <= bounds.y && visible.y + visible.height >= bounds.y + bounds.height) {
            count++;
         }
      }
      return count;
   }

   private static void showAndSelect(JTree tree, int top, int bottom, int row, boolean centerHorizontally) {
      int size = tree.getRowCount();
      if (size == 0) {
         tree.clearSelection();
         return;
      }
      if (top < 0) {
         top = 0;
      }
      if (bottom >= size) {
         bottom = size - 1;
      }
      Rectangle topBounds = tree.getRowBounds(top);
      Rectangle bottomBounds = tree.getRowBounds(bottom);
      Rectangle bounds;
      if (topBounds == null) {
         bounds = bottomBounds;
      } else if (bottomBounds == null) {
         bounds = topBounds;
      } else {
         bounds = topBounds.union(bottomBounds);
      }
      if (bounds != null) {
         TreePath path = tree.getPathForRow(row);
         if (path != null && path.getParentPath() != null) {
            Rectangle parentBounds = tree.getPathBounds(path.getParentPath());
            if (parentBounds != null) {
               bounds.x = parentBounds.x;
            }
         }
         if (!centerHorizontally) {
            bounds.x = 0;
            bounds.width = tree.getWidth();
         } else {
            bounds.width = Math.min(bounds.width, tree.getVisibleRect().width);
         }
         tree.scrollRectToVisible(bounds);
      }
      tree.setSelectionRow(row);
   }
}