The sort order is actually just controlled by permuting the order of the lines
for ii in range(shape[0]):
for jj in range(shape[1]):
for kk in range(shape[2]):
To sort as you wish, the only change needed is to instead write
for kk in range(shape[2]):
for jj in range(shape[1]):
for ii in range(shape[0]):
Don't forget to change in two places, for spin-polarized and not!
The rule is simple: the coordinate to sort primarily by, comes first (Z is kk in shape[2]), then the next-level sorting (Y, in your case, jj in shape[1]), and finally the third-level (the coordinate which changes immediately, in your case X, ii in shape[0]).
However, I also note that the print order is a bit different in your example; the script prints Z, Y, X, so if you really want the order Z, X, Y, you also must modify the print statement.