Files
gtav-src/script/dev_ng/shared/include/public/kd_tree.sch
T
2025-09-29 00:52:08 +02:00

477 lines
15 KiB
XML
Executable File

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Name: kd_tree.sch //
// Description: Implementation of multidimensional kd-trees with option to search for closest points //
// in given radius. Example use in comments below. //
// Written by: Tymon //
// Date: 2/2/2016 //
//////////////////////////////////////////////////////////////////////////////////////////////////////////////
USING "net_include.sch"
/* ***********************************************************************************************************
SAMPLE USE:
**************************************************************************************************************
_T_KD_TREE_GET_OBJECT_SORTABLE_VALUE and _T_KD_TREE_GET_SEARCH_OBJECT_VALUE are used to retrieve values of
elements we store in KD tree. By passing axis argument they serve as a "generic" implementation
of n dimensional vectors.
KD_TREE is a stuct holding all the nodes in memory. It has a pointer to function to retrieve value of each
element in each dimension.
To create the KD_TREE we need to implement _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE and create a list of numeric
indices that serves as an input for tree creation:
INT kdtreeIndices[6] // This is our list of indices
// Implementation of _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE
FUNC FLOAT GET_NUM_LIST_VALUE(INT iElementID, INT iAxis = 0)
VECTOR list[6]
list[0] = <<2, 3, 0>>
list[1] = <<5, 4, 0>>
list[2] = <<9, 6, 0>>
list[3] = <<4, 7, 0>>
list[4] = <<8, 1, 0>>
list[5] = <<7, 2, 0>>
IF iElementID >= 0 AND iElementID < 6 AND iAxis >= 0 and iAxis < 2
IF iAxis = 0
RETURN list[iElementID].x
ELSE
RETURN list[iElementID].y
ENDIF
ENDIF
RETURN 0.0
ENDFUNC
// Implementation of _T_KD_TREE_GET_SEARCH_OBJECT_VALUE - used to describe point we're quering
FUNC FLOAT GET_SEARCH_VALUE(INT iAxis = 0)
IF iAxis = 0
RETURN 3.0
ENDIF
RETURN 6.0
ENDFUNC
PROC DO_KDTREE_TEST()
kdtreeIndices[0] = 0
kdtreeIndices[1] = 1
kdtreeIndices[2] = 2
kdtreeIndices[3] = 3
kdtreeIndices[4] = 4
kdtreeIndices[5] = 5
KD_TREE kdtree = CONSTRUCT_KD_TREE(&GET_NUM_LIST_VALUE, 2, kdtreeIndices, 6)
#IF IS_DEBUG_BUILD
DEBUG_PRINT_KD_TREE(kdtree)
#ENDIF
INT iClosestPoints[10]
INT iClosestPointsCount
FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[0], &GET_SEARCH_VALUE, 3.0, iClosestPoints, iClosestPointsCount)
#IF IS_DEBUG_BUILD
INT i
PRINTLN("[KDTREE] DO_KDTREE_TEST - got search results, there's ", iClosestPointsCount, " of them: ")
NET_PRINT("[KDTREE] DO_KDTREE_TEST - ")
REPEAT iClosestPointsCount i
NET_PRINT_INT(iClosestPoints[i]) NET_PRINT(" - ")
DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, iClosestPoints[i])
IF i < iClosestPointsCount - 1
NET_PRINT(", ")
ENDIF
ENDREPEAT
#ENDIF
ENDPROC
*********************************************************************************************************** */
#IF IS_DEBUG_BUILD
PROC DEBUG_PRINT_KD_TREE_INDICES(INT &indices[], INT iFrom, INT iTo)
NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE_INDICES - Printing from ")
NET_PRINT_INT(iFrom)
NET_PRINT(" to ")
NET_PRINT_INT(iTo)
NET_PRINT(": ")
INT i
FOR i = iFrom TO iTo
NET_PRINT_INT(indices[i])
IF i < iTo
NET_PRINT(", ")
ENDIF
ENDFOR
NET_NL()
ENDPROC
PROC DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(KD_TREE &kdtree, INT iElementID)
INT i
NET_PRINT("(")
REPEAT kdtree.iDimension i
NET_PRINT_FLOAT(CALL kdtree.getValue(iElementID, i))
IF i < kdtree.iDimension - 1
NET_PRINT(", ")
ENDIF
ENDREPEAT
NET_PRINT(")")
ENDPROC
#ENDIF
PROC _SWAP_IN_KD_TREE(INT &in[], INT a, INT b)
INT iTmp = in[a]
in[a] = in[b]
in[b] = iTmp
ENDPROC
PROC _INSERT_SORT_KD_TREE(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT iAxis = 0)
INT i, j
FOR i = iFrom + 1 TO iTo
j = i
WHILE j > iFrom AND (CALL kdtree.getValue(indices[j-1], iAxis)) > (CALL kdtree.getValue(indices[j], iAxis))
_SWAP_IN_KD_TREE(indices, j, j-1)
j = j-1
ENDWHILE
ENDFOR
#IF IS_DEBUG_BUILD
PRINTLN("[KDTREE] _INSERT_SORT_KD_TREE - Sorted from ", iFrom, " to ", iTo, ". Sorted on axis ", iAxis, ". Indices after sort: ")
DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
NET_PRINT("[KDTREE] _INSERT_SORT_KD_TREE - Sorted objects: ")
FOR i = iFrom TO iTo
DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, indices[i])
IF i < iTo
NET_PRINT(", ")
ENDIF
ENDFOR
NET_NL()
#ENDIF
ENDPROC
PROC _KD_TREE_QUICK_SORT(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT iAxis = 0)
IF iFrom < iTo
INT l = iFrom + 1
INT r = iTo
INT p = indices[iFrom]
FLOAT pValue = (CALL kdtree.getValue(p, iAxis))
WHILE l < r
IF (CALL kdtree.getValue(indices[l], iAxis)) <= pValue
l = l + 1
ELIF (CALL kdtree.getValue(indices[r], iAxis)) >= pValue
r = r - 1
ELSE
_SWAP_IN_KD_TREE(indices, l, r)
ENDIF
ENDWHILE
IF (CALL kdtree.getValue(indices[l], iAxis)) < pValue
_SWAP_IN_KD_TREE(indices, l, iFrom)
l = l - 1
ELSE
l = l - 1
_SWAP_IN_KD_TREE(indices, l, iFrom)
ENDIF
_KD_TREE_QUICK_SORT(kdtree, indices, iFrom, l, iAxis)
_KD_TREE_QUICK_SORT(kdtree, indices, r, iTo, iAxis)
#IF IS_DEBUG_BUILD
NET_PRINT("[KDTREE] _KD_TREE_QUICK_SORT - Sorted from ") NET_PRINT_INT(iFrom) NET_PRINT(" to ") NET_PRINT_INT(iTo) NET_PRINT(". Sorted on axis ") NET_PRINT_INT(iAxis) NET_PRINT( ". Indices after sort: ") NET_NL()
DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
NET_PRINT("[KDTREE] _KD_TREE_QUICK_SORT - Sorted objects: ")
INT i
FOR i = iFrom TO iTo
DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, indices[i])
IF i < iTo
NET_PRINT(", ")
ENDIF
ENDFOR
NET_NL()
#ENDIF
ENDIF
ENDPROC
FUNC INT _BUILD_KD_TREE_NODE(KD_TREE &kdtree, INT &indices[], INT iFrom, INT iTo, INT depth = 0)
INT iLength = iTo - iFrom + 1
IF iLength <= 0
RETURN -1
ENDIF
PRINTLN("[KDTREE] _BUILD_KD_TREE_NODE - Building from ", iFrom, " to ", iTo, " at depth ", depth, ", length is ", iLength)
#IF IS_DEBUG_BUILD
DEBUG_PRINT_KD_TREE_INDICES(indices, iFrom, iTo)
#ENDIF
INT iAxis = depth % kdtree.iDimension
_KD_TREE_QUICK_SORT(kdtree, indices, iFrom, iTo, iAxis)
//_INSERT_SORT_KD_TREE(kdtree, iNewIndices, iFrom, iTo, iAxis)
INT iPivot = iFrom + FLOOR(TO_FLOAT(iLength) * 0.5)
PRINTLN("[KDTREE] _BUILD_KD_TREE_NODE - Choosing pivot at ", iPivot, " which is ", indices[iPivot])
INT iNodeID = kdtree.iNodesCount
kdtree.iNodesCount += 1
kdtree.nodePool[iNodeID].iElementID = indices[iPivot]
PRINTLN("[KDTREE][INSERT] _BUILD_KD_TREE_NODE - Inserted into pool at ", iNodeID, " element id ", indices[iPivot])
kdtree.nodePool[iNodeID].iLeftNodePointer = _BUILD_KD_TREE_NODE(kdtree, indices, iFrom, iPivot - 1, depth + 1)
kdtree.nodePool[iNodeID].iRightNodePointer = _BUILD_KD_TREE_NODE(kdtree, indices, iPivot + 1, iTo, depth + 1)
RETURN iNodeID
ENDFUNC
PROC CONSTRUCT_KD_TREE(KD_TREE &out, _T_KD_TREE_GET_OBJECT_SORTABLE_VALUE getValue, INT iDimension, INT &indices[], INT iLength)
PRINTLN("[KDTREE] CONSTRUCT_KD_TREE - OK let's do it, dimensions: ", iDimension, " and number of points: ", iLength)
#IF IS_DEBUG_BUILD
DEBUG_PRINT_KD_TREE_INDICES(indices, 0, iLength)
#ENDIF
out.iDimension = iDimension
out.getValue = getValue
out.iNodesCount = 0
_BUILD_KD_TREE_NODE(out, indices, 0, iLength)
ENDPROC
FUNC FLOAT _GET_SQ_DIST_OF_KD_TREE_ELEMENT_AND_SEARCH_OBJECT(KD_TREE &kdtree, KD_TREE_NODE &currentNode, _T_KD_TREE_GET_SEARCH_OBJECT_VALUE getSearchObjectValue)
INT i
FLOAT fDistSq, fTmp
REPEAT kdtree.iDimension i
fTmp = (CALL kdtree.getValue(currentNode.iElementID, i)) - (CALL getSearchObjectValue(i))
fDistSq += (fTmp * fTmp)
ENDREPEAT
RETURN fDistSq
ENDFUNC
PROC FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(KD_TREE &kdtree, KD_TREE_NODE &currentNode, _T_KD_TREE_GET_SEARCH_OBJECT_VALUE getSearchObjectValue, FLOAT fRadius, INT &out[], INT &iOutCount, INT depth = 0)
BOOL bGoingLeft
INT iAxis = depth % kdtree.iDimension
FLOAT fSearchDistSq = fRadius * fRadius
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We're at node with element ID ", currentNode.iElementID, " at the axis ", iAxis)
ENDIF
#ENDIF
FLOAT fSearchValue = CALL getSearchObjectValue(iAxis)
FLOAT fTreeValue = CALL kdtree.getValue(currentNode.iElementID, iAxis)
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - searchValue is ", fSearchValue, " while treeValue is ", fTreeValue)
ENDIF
#ENDIF
IF fSearchValue > fTreeValue
bGoingLeft = FALSE
// Go right
IF currentNode.iRightNodePointer > -1
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Going right (", fSearchValue, " > ", fTreeValue, ")")
ENDIF
#ENDIF
FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iRightNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
ELSE
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - No right to go.")
ENDIF
#ENDIF
ENDIF
ELSE
bGoingLeft = TRUE
// Go left
IF currentNode.iLeftNodePointer > -1
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Going left (", fSearchValue, " <= ", fTreeValue, ")")
ENDIF
#ENDIF
FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iLeftNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
ELSE
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - No left to go.")
ENDIF
#ENDIF
ENDIF
ENDIF
// Check if there might be any other qualifying points on the other side of the splitting plane
IF ABSF( (CALL getSearchObjectValue(iAxis)) - (CALL kdtree.getValue(currentNode.iElementID, iAxis)) ) <= fRadius
IF bGoingLeft
IF currentNode.iRightNodePointer > -1
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We are at ", currentNode.iElementID ,". There might be some more points on the right, going right.")
ENDIF
#ENDIF
FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iRightNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
ENDIF
ELSE
IF currentNode.iLeftNodePointer > -1
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - We are at ", currentNode.iElementID ,". There might be some more points on the left, going left.")
ENDIF
#ENDIF
FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS(kdtree, kdtree.nodePool[currentNode.iLeftNodePointer], getSearchObjectValue, fRadius, out, iOutCount, depth + 1)
ENDIF
ENDIF
ENDIF
FLOAT fDistSq = _GET_SQ_DIST_OF_KD_TREE_ELEMENT_AND_SEARCH_OBJECT(kdtree, currentNode, getSearchObjectValue)
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - fDistSq is ", fDistSq, " and fSearchDistSq is ", fSearchDistSq)
ENDIF
#ENDIF
IF fDistSq <= fSearchDistSq
IF iOutCount < COUNT_OF(out)
out[iOutCount] = currentNode.iElementID
iOutCount += 1
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Added current node (", currentNode.iElementID ,") to results, results count is ", iOutCount)
ENDIF
#ENDIF
ELSE
#IF IS_DEBUG_BUILD
IF g_SimpleInteriorData.bVerboseKDTree
PRINTLN("[KDTREE] FIND_ELEMENTS_IN_KD_TREE_IN_RADIUS - Reached max number of results!")
ENDIF
#ENDIF
ENDIF
ENDIF
ENDPROC
#IF IS_DEBUG_BUILD
/// PURPOSE:
/// Prints one node like: ((x, y), (node), (node))
/// PARAMS:
/// node -
/// kdtree -
/// iSpaces -
PROC _DEBUG_PRINT_KD_NODE(KD_TREE_NODE &node, KD_TREE &kdtree, INT iSpaces = 0)
INT i
REPEAT (iSpaces * 2) i
//NET_PRINT(" ")
ENDREPEAT
NET_PRINT("(")
DEBUG_PRINT_KD_TREE_SINGLE_OBJECT(kdtree, node.iElementID)
NET_PRINT(", ")
IF node.iLeftNodePointer > -1
_DEBUG_PRINT_KD_NODE(kdtree.nodePool[node.iLeftNodePointer], kdtree)
ELSE
NET_PRINT("None")
ENDIF
NET_PRINT(", ")
IF node.iRightNodePointer > -1
_DEBUG_PRINT_KD_NODE(kdtree.nodePool[node.iRightNodePointer], kdtree)
ELSE
NET_PRINT("None")
ENDIF
NET_PRINT(")")
ENDPROC
PROC DEBUG_PRINT_KD_TREE(KD_TREE &kdtree)
NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE - Node pool size: ") NET_PRINT_INT(kdtree.iNodesCount)
NET_PRINT("[KDTREE] DEBUG_PRINT_KD_TREE - ")
IF kdtree.iNodesCount > 0
_DEBUG_PRINT_KD_NODE(kdtree.nodePool[0], kdtree)
ELSE
NET_PRINT("EMPTY")
ENDIF
NET_NL()
ENDPROC
PROC DEBUG_DRAW_KD_TREE(KD_TREE &kdtree)
INT i
VECTOR vMaxCoord = <<-9999.9, -9999.9, 0.0>>
VECTOR vMinCoord = <<9999.9, 9999.9, 0.0>>
VECTOR vCurrent
REPEAT kdtree.iNodesCount i
vCurrent.X = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 0)
vCurrent.Y = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 1)
IF vCurrent.X > vMaxCoord.X
vMaxCoord.X = vCurrent.X
ENDIF
IF vCurrent.Y > vMaxCoord.Y
vMaxCoord.Y = vCurrent.Y
ENDIF
IF vCurrent.X < vMinCoord.X
vMinCoord.X = vCurrent.X
ENDIF
IF vCurrent.Y < vMinCoord.Y
vMinCoord.Y = vCurrent.Y
ENDIF
ENDREPEAT
FLOAT fNormalisedX, fNormalisedY
FLOAT fXScale = 1.0 / GET_ASPECT_RATIO(FALSE)
REPEAT kdtree.iNodesCount i
vCurrent.X = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 0)
vCurrent.Y = CALL kdtree.getValue(kdtree.nodePool[i].iElementID, 1)
fNormalisedX = (vCurrent.X - vMinCoord.X) / ABSF(vMaxCoord.X - vMinCoord.X)
fNormalisedY = (1.0 - (vCurrent.Y - vMinCoord.Y) / ABSF(vMaxCoord.Y - vMinCoord.Y))
//PRINTLN("[KDTREE] DEBUG_DRAW_KD_TREE - fNormalisedX: ", fNormalisedX, ", fNormalisedY: ", fNormalisedY, ", vMinCoord: ", vMinCoord, ", vMaxCoord: ", vMaxCoord, ", vCurrent: ", vCurrent)
DRAW_RECT( (fNormalisedX * 0.8 + 0.1) * fXScale, (fNormalisedY * 0.8 + 0.1), 0.01 * fXScale, 0.01, 255, 0, 0, 255)
ENDREPEAT
// Draw player position
vCurrent = GET_PLAYER_COORDS(PLAYER_ID())
fNormalisedX = (vCurrent.X - vMinCoord.X) / ABSF(vMaxCoord.X - vMinCoord.X)
fNormalisedY = (1.0 - (vCurrent.Y - vMinCoord.Y) / ABSF(vMaxCoord.Y - vMinCoord.Y))
FLOAT fPlayerX = (fNormalisedX * 0.8 + 0.1) * fXScale
FLOAT fPlayerY = (fNormalisedY * 0.8 + 0.1)
DRAW_RECT(fPlayerX, fPlayerY , 0.01 * fXScale, 0.01, 0, 255, 0, 255)
FLOAT fPlayerHeading = GET_ENTITY_HEADING(PLAYER_PED_ID())
DRAW_LINE_2D(fPlayerX, fPlayerY, fPlayerX + (SIN(fPlayerHeading) * 0.8 + 0.1) * fXScale, fPlayerY + (COS(fPlayerHeading) * 0.8 + 0.1), 0.002, 0, 255, 0, 255)
ENDPROC
#ENDIF