Small improvement of kd-tree