# create matrix with only rows in X corresponding to arrY if x0 or x1 is in that row
def pruneData(x0,x1, arrX, arrY):
    l = len(arrY[arrY == x0]) + len(arrY[arrY == x1])
    pArrX = np.zeros((arrX.shape[0], l))
    pArrY = np.zeros(l)
    j = 0
    for i in range(arrY.shape[0]):
        if arrY[i] == x0 or arrY[i] == x1:
            pArrX[:,j] = arrX[:,i:i+1][:,0]
            pArrY[j] = arrY[i]
            j+=1
    pArrX = np.array(pArrX)

    return pArrX, np.array(pArrY)