python: split numpy array based on values in the array,"arr[:, (1)]"