github.com/coyove/common@v0.0.0-20240403014525-f70e643f9de8/quadtree/neighbour.go (about)

     1  package quadtree
     2  
     3  import (
     4  	"bytes"
     5  	"sort"
     6  	"strings"
     7  )
     8  
     9  // http://web.archive.org/web/20120907211934/http://ww1.ucmss.com/books/LFS/CSREA2006/MSV4517.pdf
    10  func walk(code []byte, dir string, newcode []byte) {
    11  	copy(newcode, code)
    12  	for i := len(code) - 1; i >= 0; i-- {
    13  		c := code[i]
    14  		newcode[i], dir = walkFSM(c, dir)
    15  		if dir == "" {
    16  			break
    17  		}
    18  	}
    19  }
    20  
    21  func walkFSM(startOrth byte, dir string) (byte, string) {
    22  	walkOrth := func(startOrth byte, dir byte) (byte, byte) {
    23  		// 1 0 1 0
    24  		// 2 3 2 3
    25  		// 1 0 1 0
    26  		// 2 3 2 3
    27  		switch dir {
    28  		case 'u':
    29  			return ("\x03\x02\x01\x00"[startOrth]), "uu\x00\x00"[startOrth]
    30  		case 'd':
    31  			return ("\x03\x02\x01\x00"[startOrth]), "\x00\x00dd"[startOrth]
    32  		case 'l':
    33  			return ("\x01\x00\x03\x02"[startOrth]), "\x00ll\x00"[startOrth]
    34  		case 'r':
    35  			return ("\x01\x00\x03\x02"[startOrth]), "r\x00\x00r"[startOrth]
    36  		default:
    37  			panic(dir)
    38  		}
    39  	}
    40  	switch dir {
    41  	case "u", "d", "l", "r":
    42  		o, d := walkOrth(startOrth, dir[0])
    43  		return o, strings.Trim(string(d), "\x00")
    44  	case "ul", "ur", "dl", "dr":
    45  		o, d := walkOrth(startOrth, dir[0])
    46  		o2, d2 := walkOrth(o, dir[1])
    47  		return o2, strings.Trim(string(d)+string(d2), "\x00")
    48  	default:
    49  		panic(dir)
    50  	}
    51  }
    52  
    53  func (t QuadTree) getByCode(locode []byte) (interface{}, error) {
    54  	var err error
    55  	for _, code := range locode {
    56  		if t.isleaf() {
    57  			return t.Elems, nil
    58  		}
    59  		o := t.O[code]
    60  		if o == "" {
    61  			return nil, nil
    62  		}
    63  		t, err = t.LoadTree(o)
    64  		if err != nil {
    65  			return nil, err
    66  		}
    67  	}
    68  	// d, n := t.MaxDepth()
    69  	// fmt.Println("####code:", locode, d, n)
    70  	return t, nil // return the parent tree whose children all share the same code prefix
    71  }
    72  
    73  func (t QuadTree) FindNeig(src Point, distance func(p Point) float64) ([]Element, error) {
    74  	if distance == nil { // Simple Euclidean distance
    75  		distance = src.Distance
    76  	}
    77  
    78  	buf := &bytes.Buffer{}
    79  	cands := map[Point]Element{}
    80  
    81  	_, tid, _ := t.find(buf, src)
    82  	pt, err := t.LoadTree(tid)
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	for _, e := range pt.Elems {
    87  		if e.Point != src {
    88  			cands[e.Point] = e
    89  		}
    90  	}
    91  
    92  	x, tmp := buf.Bytes(), make([]byte, buf.Len()*8)
    93  NEXT_DIR:
    94  	for i, dir := range directions {
    95  		y := tmp[i*len(x) : (i+1)*len(x)]
    96  		walk(x, dir, y) // walk dir from x to y
    97  		for ii := 0; ii < i; ii++ {
    98  			if bytes.Equal(tmp[ii*len(x):(ii+1)*len(x)], y) {
    99  				continue NEXT_DIR
   100  			}
   101  		}
   102  		v, err := t.getByCode(y)
   103  		if err != nil {
   104  			return nil, err
   105  		}
   106  		switch v := v.(type) {
   107  		case map[Point]Element:
   108  			for p, e := range v {
   109  				cands[p] = e
   110  			}
   111  		case QuadTree:
   112  			v.Iterate(func(e Element) error { cands[e.Point] = e; return nil })
   113  		}
   114  	}
   115  	els := make([]Element, 0, len(cands))
   116  	for p, k := range cands {
   117  		if p != (Point{}) {
   118  			els = append(els, k)
   119  		}
   120  	}
   121  	sort.Slice(els, func(i, j int) bool { return distance(els[i].Point) < distance(els[j].Point) })
   122  	return els, nil
   123  }