一文讲懂GeoHash（二）

本文是一文讲动GeoHash的第二篇，代码实现篇，主要是代大家从0到1实现GeoHash，代码量不多，大约只有400多行，如果还不太清楚GeoHash的原理，请参考[一文讲懂GeoHash（一）]

首先是关于存储哈希的数据结构，我并没有使用redis的GeoHash所使用的跳表，而是使用了字典树Trie，因为跳表的时间复杂度主要取决于数据量，而Trie的时间复杂度主要取决与字符串长度，所以显然用字典树能尽可能提高GeoHash的性能。
而之所以redis使用跳表作为GeoHash的底层数据结构，我想很大的一个原因就是为了与ZSet复用，因为ZSet的底层也是用的跳表。

再说一下GeoHash的架构：
• GEOService：geohash 服务类，用户访问的统一入口.

• mu：通过一把读写锁保护 geohash 前缀树的并发安全性

• root：持有的 geohash 前缀树的根节点引用

• geoTrieNode：geohash 前缀树中的节点类型

• children：孩子节点列表. 由于采用 base32 编码，因此最多只能有 32 个 child . 其中 children[0] 对应值为 ‘0’ 的 child，children[31] 对应值为 ‘z’ 的 child，以此类推

• passCnt：途径该节点的子孙节点计数器

• end：标志是否有字符串以当前节点结尾

• GEOEntry：对应与一个 geohash 字符串的矩形区域

• Points： 矩形区域内的点集合

• Hash：geohash 字符串

• GEOPoint：通过经、纬度标识的一个点

• Lng、Lat：经纬度

• Val：点对应的业务信息

至于GeoHash的代码，我的每一步都作了详细注释，大家可以跟着动手实现一下

package main

import (
	"fmt"
	"math"
	"strconv"
	"strings"
	"sync"
)

type GeoService struct {
	mu   sync.RWMutex //控制并发
	root *geotrie     //前缀树根节点
}

type geotrie struct {
	childern  [32]*geotrie //子孙节点
	passCnt   int          //经过该节点的数量
	end       bool         //是否有以该节点为结尾
	GeoDetail              //区域信息
}

type GeoDetail struct {
	Points   map[string]interface{} //区域内点的集合
	hashcode string                 //区域哈希编码
}

type GeoPoints struct {
	longitude float64     //经度
	latitude  float64     //纬度
	Val       interface{} //
}

// 创建一个新的GEO服务
func NewGeoservice() *GeoService {
	return &GeoService{
		root: &geotrie{},
	}
}

// 获取所有节点
func (g *GeoDetail) GetAllPoints() []*GeoPoints {
	points := make([]*GeoPoints, 0, len(g.Points))
	for key, val := range g.Points {
		lng, lat := g.lngLat(key)
		points = append(points, &GeoPoints{
			longitude: lng,
			latitude:  lat,
			Val:       val,
		})
	}
	return points
}

// 添加节点
func (g *GeoDetail) AddPoints(lng, lat float64, val interface{}) {
	if g.Points == nil {
		g.Points = map[string]interface{}{}
	}
}

func (g *GeoDetail) pointKey(lng, lat float64) string {
	return fmt.Sprintf("%v_%v", lng, lat)
}

// 将 point 对应的 key ${lng}_${lat} 转为对应的经纬度
func (g *GeoDetail) lngLat(key string) (lng float64, lat float64) {
	info := strings.Split(key, "_")
	lng, _ = strconv.ParseFloat(info[0], 64)
	lat, _ = strconv.ParseFloat(info[1], 64)
	return
}

// 将十进制数值转为 base32 编码
var Base32 = []byte{'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'M', 'N', 'P',
	'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'}

// 将 base32 编码转为十进制数值
func (g *GeoService) base32ToIndex(base32 byte) int {
	if base32 >= '0' && base32 <= '9' {
		return int(base32 - '0')
	}
	if base32 >= 'B' && base32 <= 'H' {
		return int(base32 - 'B' + 26)
	}
	if base32 >= 'J' && base32 <= 'K' {
		return int(base32 - 'J' + 33)
	}
	if base32 >= 'M' && base32 <= 'N' {
		return int(base32 - 'J' + 35)
	}
	if base32 >= 'P' && base32 <= 'Z' {
		return int(base32 - 'J' + 37)
	}
	return -1
}

func (g *GeoService) binarybitsresult(bits strings.Builder, val, start, end float64) string {
	mid := (end + start) / 2
	if val <= mid {
		bits.WriteByte(0)
		end = mid
	} else {
		bits.WriteByte(1)
		start = mid
	}
	if len(bits.String()) >= 20 {
		return bits.String()
	}

	return g.binarybitsresult(bits, val, start, end)
}

// 将经纬度转化为字符串
func (g *GeoService) Hash(lng, lat float64) string {
	//二分经度
	tlng := g.binarybitsresult(strings.Builder{}, lng, -180, 180)
	//二分纬度
	tlat := g.binarybitsresult(strings.Builder{}, lat, -90, 90)
	//将经度和纬度交互5个一组合并
	var FiveBits strings.Builder
	var realBits strings.Builder
	for i := 1; i <= 40; i++ {
		if i&1 == 0 {
			FiveBits.WriteByte(tlng[(i-1)>>1])
		} else {
			FiveBits.WriteByte(tlat[(i-1)>>1])
		}
		if i%5 != 0 {
			continue
		}
		val, _ := strconv.ParseInt(FiveBits.String(), 2, 64)
		realBits.WriteByte(Base32[val])
		FiveBits.Reset()
	}
	//返回结果
	return realBits.String()
}

// 查询结果
func (g *GeoService) Get(hashcode string) (*GeoDetail, bool) {
	//加读锁
	g.mu.RLock()
	defer g.mu.RUnlock()

	//获取结果
	target := g.get(hashcode)
	//判断是否存在且end为true
	if target == nil || !target.end {
		return nil, false
	}
	//返回结果
	return &target.GeoDetail, true
}

func (g *GeoService) get(hashcoed string) *geotrie {
	//从根节点开始遍历
	move := g.root
	//开始便利hashcode的每一个字节
	for i := 0; i < len(hashcoed); i++ {
		index := g.base32ToIndex(hashcoed[i])
		//如果索引不存在或者不合理，直接返回空即可
		if index == -1 || move.childern[index] == nil {
			return nil
		}
		//继续遍历下一个子节点
		move = move.childern[index]
	}
	return move
}

// 添加节点
func (g *GeoService) Add(lng, lat float64, val interface{}) {
	//加写锁
	g.mu.Lock()
	defer g.mu.Unlock()
	//获取哈希值
	hashcode := g.Hash(lng, lat)

	target := g.get(hashcode)
	//如果target不为空且end标识符为true，则直接插入
	if target != nil && target.end {
		target.AddPoints(lng, lat, val)
		return
	}

	//如果不存在，则从头开始遍历
	move := g.root
	for i := 0; i < len(hashcode); i++ {
		index := g.base32ToIndex(hashcode[i])
		//如果不存在，则直接构建路径
		if move.childern[index] == nil {
			move.childern[index] = &geotrie{}
		}
		//路径数加一
		move.passCnt++
		move = move.childern[index]
	}
	//结束标志为真
	move.end = true
	move.AddPoints(lng, lat, val)
	move.hashcode = hashcode
}

// 前缀查询，根据前缀查询出所有结果
func (g *GeoService) Getbyprefix(prefix string) []*GeoDetail {
	//加锁
	g.mu.RLock()
	defer g.mu.RUnlock()
	target := g.get(prefix)
	//如果为空则一定不存在以此为前缀的,直接返回
	if target == nil {
		return nil
	}

	return target.getbyprefix()
}

func (ge *geotrie) getbyprefix() []*GeoDetail {
	var geodatils []*GeoDetail
	//如果为尾节点，直接添加
	if ge.end == true {
		geodatils = append(geodatils, &ge.GeoDetail)
	}
	//遍历所有非空的子结点
	for i := 0; i < len(ge.childern); i++ {
		if ge.childern[i] != nil {
			geodatils = append(geodatils, ge.childern[i].getbyprefix()...)
		}
	}
	return geodatils
}

// 删除节点
func (g *GeoService) del(hashcode string) bool {
	target := g.get(hashcode)
	//如果为空，则没有该节点，直接返回
	if target == nil {
		return false
	}

	//说明有该节点，从头遍历其所有经过的节点
	move := g.root
	for i := 0; i <= len(hashcode); i++ {
		index := g.base32ToIndex(hashcode[i])
		move.childern[index].passCnt--
		//如果此时减完为0,则节点直接置空
		if move.passCnt == 0 {
			move.childern[index] = nil
			return true
		}
		move = move.childern[index]
	}
	move.end = false
	return true
}

// 范围查询
func (g *GeoService) getBitsLengthByRadiusM(radiusM int) (int, error) {
	if radiusM > 160*1000 || radiusM < 0 {
		return -1, fmt.Errorf("invalid radius: %d", radiusM)
	}

	var i int
	for {
		if radiusM <= distRank[i+1] {
			return 8 - i, nil
		}
		i++
	}
}

var distRank = []int{0, 20, 150, 600, 5000, 20000, 160000}

func (g *GeoService) calDistance(lng1, lat1, lng2, lat2 float64) float64 {
	return 111 * (math.Pow(lng1-lng2, 2) + math.Pow(lat1-lat2, 2))
}

func (g *GeoService) getCenterPoints(lng, lat float64, radiusM int) [9][2]float64 {
	dif := float64(radiusM) / 111
	left := lng - dif
	if left < -180 {
		left += 360
	}
	right := lng + dif
	if right > 180 {
		right -= 360
	}
	bot := lat - dif
	if bot < -90 {
		bot += 180
	}
	top := lat + dif
	if top > 90 {
		top -= 180
	}

	return [9][2]float64{
		{
			left, top,
		},
		{
			lng, top,
		},
		{
			right, top,
		},
		{
			left, lat,
		},
		{
			lng, lat,
		},
		{
			right, lat,
		},
		{
			left, bot,
		},
		{
			lng, bot,
		},
		{
			right, bot,
		},
	}
}

// 查询指定范围范围内
// radius 单位为 m.
// m 转为对应的经纬度 经度1度≈111m；纬度1度≈111m
func (g *GeoService) ListByRadiusM(lng, lat float64, radiusM int) ([]*GeoPoints, error) {
	// 加一把读锁
	g.mu.RLock()
	defer g.mu.RUnlock()

	// 1 根据用户指定的查询范围，确定所需要的 geohash 字符串的长度，保证对应的矩形区域长度大于等于 radiusM
	bitsLen, err := g.getBitsLengthByRadiusM(2 * radiusM)
	if err != nil {
		return nil, err
	}

	// 2 针对于传入的 lng、lat 中心点，沿着上、下、左、右方向进行偏移，获取到包含自身在内的总共 9 个点的点矩阵
	// 核心是为了保证通过 9 个点获取到的矩形区域一定能完全把检索范围包含在内
	points := g.getCenterPoints(lng, lat, radiusM)

	// 3. 针对这9个点，通过 ListByPrefix 方法，分别取出区域内的所有子 GEOEntry
	var rawEntries []*GeoDetail
	for i := 0; i < len(points); i++ {
		geoHash := g.Hash(points[i][0], points[i][1])[:bitsLen]
		rawEntries = append(rawEntries, g.Getbyprefix(geoHash)...)
	}

	// 4. 针对所有 entry，取出其中包含的所有 point
	// 取出 point 之后，计算其与 center point 的相对距离，如果超过范围则进行过滤
	var geoPoints []*GeoPoints
	// 遍历所有的 entry
	for _, rawEntry := range rawEntries {
		// 遍历每个 entry 中所有的 point
		for _, rawPoint := range rawEntry.GetAllPoints() {
			// 计算一个 point 与中心点 lng、lat 的相对距离
			dist := g.calDistance(lng, lat, rawPoint.longitude, rawPoint.latitude)
			// 如果相对距离大于 radiusM，则进行过滤
			if dist > float64(radiusM) {
				continue
			}
			// 相对距离满足条件，则将 point 追加到 list 中
			geoPoints = append(geoPoints, rawPoint)
		}
	}

	return geoPoints, nil
}

总体难度并不是很大，大家有任何对代码或者思路的疑问都可以随时私信或者在评论区评论

下篇：一文读懂一致性哈希