code: public static int[] diag(int[][] a) {
      int[] rv = new int[a.length];
      for (int i = 0; i < a.length; i++) {
        rv[i] = a[i][i];
      }
      return rv;
    }
docs: Returns diagonal elements of the given (square) matrix.