
    Wi8}                        d dl mZ d dlZd dlZd dlZd dlmZm	Z	m
Z
mZmZ d dlmZmZ d dlZd dlmZ dZ ej(                  ej*                        j,                  dz   Z ej(                  ej*                        j0                  dz
  Z edg d	      Zej6                  ddd   Zej:                  dddddf   Zej>                  ddd   Z ej:                  Z! ejD                   ej*                  d
       ej*                  d
      f      Z#ej*                  ddd   Z$ ejJ                   e&d      D  cg c]  }  e'|       jQ                  d       c} ej6                        Z) ejT                   ejV                  jY                  ej*                  ddd   ej*                  ddd   ee!f      ej6                  dddddf   ej*                  ddd   ejZ                  ddd         ej\                  ej\                  ej6                  ddd   ej6                  ej6                  ej\                  ej\                  ej\                  ej\                  d	ddd      d        Z/ ejT                   ejV                  jY                  ej*                  ddd   ej*                  ddd   e e!f      ej>                  dddddf   ej*                  ddd   ejZ                  ddd         ej\                  ej\                  ej>                  ddd   ej6                  ej6                  ej\                  ej\                  ej\                  ej\                  d	ddd      d        Z0 ejT                   ejV                  jY                  ej*                  ddd   ej*                  ddd   ee!f      ej6                  dddddf   ej*                  ddd   ejZ                  ddd         ej\                  ej\                  ej6                  ddd   ej6                  ej6                  ej\                  ej\                  ej\                  ej\                  d	ddd      d        Z1 ejT                  dddejV                  j6                  ddd   ejV                  j6                  ddd   ejV                  j6                  ejV                  j\                  d      d        Z2 ejT                  ddd      d        Z3dZ4 ejT                  ddd      d        Z5 ejT                  ddd      d        Z6 ejT                  ddd      dad       Z7dZ8 ejT                  ddd      d        Z9 ejT                  ddd      d         Z: ejT                  ddejV                  j*                  ejV                  j*                  d!"      	 	 dbd#       Z; ejT                  ddejV                  j*                  ejV                  j*                  d!"      	 	 dbd$       Z< ejT                  dd%      	 	 	 dcd&       Z= ejT                  ddd      d'        Z> ejT                  ddd      d(        Z? ejT                  ddejV                  j*                  ejV                  j*                  d!"      	 	 dbd)       Z@ ejT                  ddejV                  j*                  ejV                  j*                  d!"      	 	 dbd*       ZA ejT                  dd%      	 	 	 dcd+       ZB ejT                  ddd      ddd-       ZC ejT                  ddd      d.        ZD ejT                  ddejV                  j*                  ejV                  j*                  d!"      	 	 dbd/       ZE ejT                  dd%      	 	 dbd0       ZF ejT                  dejV                  j*                  ejV                  j*                  d!1      	 	 dbd2       ZG ejT                  dejV                  j                  e#      ejV                  j*                  ejV                  j*                  d31      	 	 dbd4       ZI ejT                  dejV                  j                  e#      ejV                  j*                  ejV                  j*                  d31      	 	 dbd5       ZJ ejT                  dejV                  j*                  ejV                  j*                  d!1      	 	 dbd6       ZK ejT                  dejV                  j*                  ejV                  j*                  d!1      	 	 dbd7       ZL ejT                  d8      dcd9       ZM ejT                  d8      	 	 	 dcd:       ZN ejT                  d8      dcd;       ZO ejT                  d<ejV                  j                  ejV                  j                  ejV                  j6                  dd=d>      ejV                  j6                  ejV                  j                  ejV                  j6                  dd=d>      ejV                  j                  ejV                  jZ                  dd=d>            gdejV                  j6                  ejV                  j                  ejV                  j                  d?d@      dA        ZT ejT                  dBejV                  j                  ejV                  j                  ejV                  j>                  dd=d>      ejV                  j6                  ejV                  j                  ejV                  j>                  dd=d>      ejV                  j                  ejV                  jZ                  dd=d>            gdejV                  j6                  ejV                  j                  ejV                  j                  d?d@      dC        ZU ejT                  dD ejV                  j                  ejV                  j*                  dd=d>      ejV                  j                  ejV                  j6                  dd=d>      ejV                  j                  ejV                  j6                  dEd=d>      ejV                  j                  ejV                  j6                  dd=d>      ejV                  j                  ejV                  j*                  dEd=d>      ejV                  j                  ejV                  j*                  dd=d>      ejV                  j                  ejV                  jZ                  dd=d>            gejV                  j\                  ejV                  j                  dFdG      dH        ZV ejT                  dI ejV                  j                  ejV                  j*                  dd=d>      ejV                  j                  ejV                  j>                  dd=d>      ejV                  j                  ejV                  j>                  dEd=d>      ejV                  j                  ejV                  j6                  dd=d>      ejV                  j                  ejV                  j*                  dEd=d>      ejV                  j                  ejV                  j*                  dd=d>      ejV                  j                  ejV                  jZ                  dd=d>            gejV                  j\                  ejV                  j                  dFdG      dJ        ZW ejT                  ddK      dL        ZX ejT                  dMejV                  j\                  idG      dN        ZY	 	 	 	 dedOZZ ejT                  d8      dP        Z[dQ Z\dR Z]dS Z^ ejT                         dT        Z_ ejT                  dU      dV        Z`dW Zad ZbdZcdEZdd,ZedXZfdY ZgdZ Zh ejT                  dejZ                  ddd   ej6                  ej\                  d[d\      d]        Zi ejT                  dej*                  ej*                  d^d_      d`        Zjyc c} w )f    )warnN)
sparse_mulsparse_diff
sparse_sumarr_intersectsparse_dot_product)tau_rand_intnorm)
namedtupleg:0yE>   FlatTreehyperplanesoffsetschildrenindices	leaf_size   1dtype)	n_leftn_righthyperplane_vectorhyperplane_offsetmargindi
left_indexright_indexT)localsfastmathnogilcachec                    | j                   d   }t        |      |j                   d   z  }t        |      |j                   d   z  }|||k(  z  }||j                   d   z  }||   }||   }t        | |         }t        | |         }	t        |      t        k  rd}t        |	      t        k  rd}	t        j                  |t
        j                        }
t        |      D ]  }| ||f   |z  | ||f   |	z  z
  |
|<    t        |
      }t        |      t        k  rd}t        |      D ]  }|
|   |z  |
|<    d}d}t        j                  |j                   d   t
        j                        }t        |j                   d         D ]x  }d}t        |      D ]  }||
|   | ||   |f   z  z  } t        |      t        k  r%t        |      dz  ||<   ||   dk(  r|dz  }Y|dz  }_|dkD  rd||<   |dz  }od||<   |dz  }z |dk(  s|dk(  rEd}d}t        |j                   d         D ]&  }t        |      dz  ||<   ||   dk(  r|dz  }"|dz  }( t        j                  |t
        j                        }t        j                  |t
        j                        }d}d}t        |j                   d         D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' |||
dfS )M  Given a set of ``graph_indices`` for graph_data points from ``graph_data``, create
    a random hyperplane to split the graph_data, returning two arrays graph_indices
    that fall on either side of the hyperplane. This is the basis for a
    random projection tree, which simply uses this splitting recursively.
    This particular split uses cosine distance to determine the hyperplane
    and which side each graph_data sample falls on.
    Parameters
    ----------
    data: array of shape (n_samples, n_features)
        The original graph_data to be split
    indices: array of shape (tree_node_size,)
        The graph_indices of the elements in the ``graph_data`` array that are to
        be split in the current operation.
    rng_state: array of int64, shape (3,)
        The internal state of the rng
    Returns
    -------
    indices_left: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    indices_right: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    r   r         ?r              )shaper	   r
   absEPSnpemptyfloat32rangeint8int32)datar   	rng_statedimr    r!   leftright	left_norm
right_normr   r   hyperplane_normr   r   sider   r   indices_leftindices_rights                       `/home/sietch6/trending-topics-pipeline/venv/lib/python3.12/site-packages/pynndescent/rp_trees.pyangular_random_projection_splitr@   )   sF   X **Q-C i(7==+;;Jy)GMM!,<<K:,,Ka 00K:DK ET$Z Id5k"J
9~	
:
 BJJ73Z 
 $T1W	 9NZ' 
!

 ,-O
?c!3Z F03oE!F FG88GMM!$bgg.D7==#$ s 	AA'*T'!*a--@@@F	A v;"9-1DGAw!|!1aZDGaKFDGqLG!( {glw}}Q'( 	A"9-1DGAw!|!1	 88F"((3LHHWBHH5M FG4::a=! 7a<#*1:L aKF%,QZM'"qLG (93>>    c                    | j                   d   }t        |      |j                   d   z  }t        |      |j                   d   z  }|||k(  z  }||j                   d   z  }||   }||   }d}d}	t        j                  |dz  t        j                        }
|
d| }|
|d }t        |      D ]+  }| ||f   | ||f   z  }|| ||f   z  ||<   || ||f   z  ||<   - d}t        |      D ]3  }|t        |
|      z  }|t        | ||f      z  }|	t        | ||f      z  }	5 d}d}t        j                  |j                   d   t        j                        }t        |j                   d         D ]  }d}t        |      D ]6  }|t        ||   | ||   |f   z     z  }|t        ||   | ||   |f   z     z  }8 t        |      t        k  r%t        |      dz  ||<   ||   dk(  r|dz  }z|dz  }|dkD  rd||<   |dz  }d||<   |dz  } |dk(  s|dk(  rEd}d}t        |j                   d         D ]&  }t        |      dz  ||<   ||   dk(  r|dz  }"|dz  }( t        j                  |t        j                        }t        j                  |t        j                        }d}d}t        |j                   d         D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' |||
dfS )r'   r   r   r)   r*   r   N)r+   r	   r.   r/   uint8r1   popcntr2   r,   r-   r3   )r4   r   r5   r6   r    r!   r7   r8   r9   r:   r   positive_hyperplane_componentnegative_hyperplane_componentr   
xor_vectorr;   r   r   r<   r   r   r=   r>   s                          r?   )angular_bitpacked_random_projection_splitrH      s   X **Q-C i(7==+;;Jy)GMM!,<<K:,,Ka 00K:DK EIJ q9$5ds$;!$5cd$;!3Z I47mUAX7
+5dAg+G%a(+5eQh+H%a(I
 O3Z -6"3A"677VDqM**	fT%(^,,
- FG88GMM!$bgg.D7==#$ s 	UAf:1=WQZQR]@SSTTFf:1=WQZQR]@SSTTF	U v;"9-1DGAw!|!1aZDGaKFDGqLG#* {glw}}Q'( 	A"9-1DGAw!|!1	 88F"((3LHHWBHH5M FG4::a=! 7a<#*1:L aKF%,QZM'"qLG (93>>rA   c                    | j                   d   }t        |      |j                   d   z  }t        |      |j                   d   z  }|||k(  z  }||j                   d   z  }||   }||   }d}t        j                  |t        j                        }	t        |      D ]/  }
| ||
f   | ||
f   z
  |	|
<   ||	|
   | ||
f   | ||
f   z   z  dz  z  }1 d}d}t        j                  |j                   d   t        j                        }t        |j                   d         D ]  }|}t        |      D ]  }
||	|
   | ||   |
f   z  z  } t        |      t        k  r.t        t        |            dz  ||<   ||   dk(  r|dz  }b|dz  }h|dkD  rd||<   |dz  }xd||<   |dz  } |dk(  s|dk(  rEd}d}t        |j                   d         D ]&  }t        |      dz  ||<   ||   dk(  r|dz  }"|dz  }( t        j                  |t        j                        }t        j                  |t        j                        }d}d}t        |j                   d         D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' |||	|fS )aP  Given a set of ``graph_indices`` for graph_data points from ``graph_data``, create
    a random hyperplane to split the graph_data, returning two arrays graph_indices
    that fall on either side of the hyperplane. This is the basis for a
    random projection tree, which simply uses this splitting recursively.
    This particular split uses euclidean distance to determine the hyperplane
    and which side each graph_data sample falls on.
    Parameters
    ----------
    data: array of shape (n_samples, n_features)
        The original graph_data to be split
    indices: array of shape (tree_node_size,)
        The graph_indices of the elements in the ``graph_data`` array that are to
        be split in the current operation.
    rng_state: array of int64, shape (3,)
        The internal state of the rng
    Returns
    -------
    indices_left: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    indices_right: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    r   r   r)   r          @r*   )
r+   r	   r.   r/   r0   r1   r2   r,   r-   r3   )r4   r   r5   r6   r    r!   r7   r8   r   r   r   r   r   r<   r   r   r=   r>   s                     r?   !euclidean_random_projection_splitrK   0  s   X **Q-C i(7==+;;Jy)GMM!,<<K:,,Ka 00K:DK E BJJ73Z 
#D!G}tE1H~=!a DqMDN$BCcI	

 FG88GMM!$bgg.D7==#$ "s 	AA'*T'!*a--@@@F	A v;,y12Q6DGAw!|!1aZDGaKFDGqLG!( {glw}}Q'( 	A"9-1DGAw!|!1	 88F"((3LHHWBHH5M FG4::a=! 7a<#*1:L aKF%,QZM'"qLG (9;LLLrA   )normalized_left_datanormalized_right_datar;   r   )r#   r$   r%   r"   c                    t        |      |j                  d   z  }t        |      |j                  d   z  }|||k(  z  }||j                  d   z  }||   }||   }| ||   ||dz       }	|||   ||dz       }
| ||   ||dz       }|||   ||dz       }t        |
      }t        |      }t        |      t        k  rd}t        |      t        k  rd}|
|z  j                  t        j                        }||z  j                  t        j                        }t        |	|||      \  }}t        |      }t        |      t        k  rd}t        |j                  d         D ]  }||   |z  ||<    d}d}t        j                  |j                  d   t        j                        }t        |j                  d         D ]  }d}| |||      |||   dz       }||||      |||   dz       }t        ||||      \  }}|D ]  }||z  }	 t        |      t        k  r%t        |      dz  ||<   ||   dk(  r|dz  }{|dz  }|dkD  rd||<   |dz  }d||<   |dz  } |dk(  s|dk(  rEd}d}t        |j                  d         D ]&  }t        |      dz  ||<   ||   dk(  r|dz  }"|dz  }( t        j                  |t        j                        }t        j                  |t        j                        } d}d}t        |j                  d         D ]%  }||   dk(  r||   ||<   |dz  }||   | |<   |dz  }' t        j                  ||f      }!|| |!dfS )  Given a set of ``graph_indices`` for graph_data points from a sparse graph_data set
    presented in csr sparse format as inds, graph_indptr and graph_data, create
    a random hyperplane to split the graph_data, returning two arrays graph_indices
    that fall on either side of the hyperplane. This is the basis for a
    random projection tree, which simply uses this splitting recursively.
    This particular split uses cosine distance to determine the hyperplane
    and which side each graph_data sample falls on.
    Parameters
    ----------
    inds: array
        CSR format index array of the matrix
    indptr: array
        CSR format index pointer array of the matrix
    data: array
        CSR format graph_data array of the matrix
    indices: array of shape (tree_node_size,)
        The graph_indices of the elements in the ``graph_data`` array that are to
        be split in the current operation.
    rng_state: array of int64, shape (3,)
        The internal state of the rng
    Returns
    -------
    indices_left: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    indices_right: array
        The elements of ``graph_indices`` that fall on the "left" side of the
        random hyperplane.
    r   r   r(   r)   r*   r   )r	   r+   r
   r,   r-   astyper.   r0   r   r1   r/   r2   r   r3   vstack)"indsindptrr4   r   r5   r    r!   r7   r8   	left_inds	left_data
right_inds
right_datar9   r:   rL   rM   hyperplane_indshyperplane_datar;   r   r   r   r<   r   r   i_indsi_data_mul_datavalr=   r>   
hyperplanes"                                     r?   &sparse_angular_random_projection_splitr`     s   T i(7==+;;Jy)GMM!,<<K:,,Ka 00K:DK EVD\F4!8$45IVD\F4!8$45IfUmfUQY&78JfUmfUQY&78JYIj!J
9~	
:
 &	199"**E'*4<<RZZH'2'5J($O_ ?+O
?c!?((+, B,Q//AB FG88GMM!$bgg.D7==#$ fWQZ(6'!*q.+ABfWQZ(6'!*q.+AB /66R8 	CcMF	 v;"9-1DGAw!|!1aZDGaKFDGqLG+2 {glw}}Q'( 	A"9-1DGAw!|!1	 88F"((3LHHWBHH5M FG4::a=! 7a<#*1:L aKF%,QZM'"qLG O_=>J
C77rA   )r#   r$   r%   c                 $   t        j                  t        |            |j                  d   z  }t        j                  t        |            |j                  d   z  }|||k(  z  }||j                  d   z  }||   }||   }| ||   ||dz       }	|||   ||dz       }
| ||   ||dz       }|||   ||dz       }d}t	        |	|
||      \  }}t        |	|
||      \  }}|dz  }t        ||||j                  t         j                              \  }}|D ]  }||z  }	 d}d}t        j                  |j                  d   t         j                        }t        |j                  d         D ]  }|}| |||      |||   dz       }||||      |||   dz       }t        ||||      \  }}|D ]  }||z  }	 t        |      t        k  r.t        t        |            dz  ||<   ||   dk(  r|dz  }|dz  }|dkD  rd||<   |dz  }d||<   |dz  } |dk(  s|dk(  rNd}d}t        |j                  d         D ]/  }t        t        |            dz  ||<   ||   dk(  r|dz  }+|dz  }1 t        j                  |t         j                        }t        j                  |t         j                        }d}d}t        |j                  d         D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' t        j                  ||f      }||||fS )rO   r   r   r)   rJ   r*   r   )r.   r,   r	   r+   r   r   r   rP   r0   r/   r2   r1   r-   r3   rQ   )rR   rS   r4   r   r5   r    r!   r7   r8   rT   rU   rV   rW   r   rX   rY   offset_indsoffset_datar^   r   r   r<   r   r   rZ   r[   r\   r]   r=   r>   r_   s                                  r?   (sparse_euclidean_random_projection_splitrd   1  s   @ Y/07==3CCJ&&i01GMM!4DDK:,,Ka 00K:DK EVD\F4!8$45IVD\F4!8$45IfUmfUQY&78JfUmfUQY&78J '29j*($O_  *)Y
JWK#K)+{7I7I"**7U K  !S ! FG88GMM!$bgg.D7==#$ "fWQZ(6'!*q.+ABfWQZ(6'!*q.+AB /66R8 	CcMF	 v;,y12Q6DGAw!|!1aZDGaKFDGqLG)0 {glw}}Q'( 	A,y12Q6DGAw!|!1	 88F"((3LHHWBHH5M FG4::a=! 7a<#*1:L aKF%,QZM'"qLG O_=>J
4EEErA   i  Fc                     d}| j                   d   dz
  }||k  r+||z   dz  }| |   |k(  r|S | |   |k  r|dz   }n|dz
  }||k  r+y)z5Binary search returning index if found, -1 otherwise.r   r   r*   r   )r+   )
sorted_arrvaluelohimids        r?   binary_searchrk     so     
B			!	q	 B
(Bw1nc?e#J_u$qBqB ( rA   c                    | j                   d   }t        j                  |t        j                        }t	        |      D ]?  }t	        | j                   d         D ]"  }| ||f   }|dk\  s||k  s||xx   dz  cc<   $ A |S )a  Compute global in-degree for all points in the graph.

    In-degree of a point is how many times it appears as a neighbor of other points.
    This is computed once and reused throughout tree construction.

    Parameters
    ----------
    neighbor_indices : array of shape (n_samples, n_neighbors)
        The neighbor graph indices.

    Returns
    -------
    global_degrees : array of shape (n_samples,)
        The in-degree of each point.
    r   r   r   )r+   r.   zerosr3   r1   )neighbor_indicesn_pointsglobal_degreesr   jneighbors         r?   compute_global_degreesrs     s    *  %%a(HXXhbhh7N8_ .'--a01 	.A'1-H1}H!4x(A-(	.. rA   c                    | j                   d   }t        ||      }t        j                  |t        j                  d      t        j                        }t        j
                  |t        j                        }t        |      D ]y  }|| |      }|||dz
     kD  s|dz
  }	|	dkD  r!|||	dz
     kD  r|	dz  }	|	dkD  r|||	dz
     kD  rt        |dz
  |	d      D ]  }
||
dz
     ||
<   ||
dz
     ||
<    |||	<   | |   ||	<   { |S )a.  Get the indices of the top k highest-degree points from a subset.

    Uses an efficient O(n) selection for small k by maintaining a min-heap of k elements.

    Parameters
    ----------
    indices : array of shape (n,)
        The point indices in the current split.
    global_degrees : array of shape (n_total,)
        Precomputed global degrees for all points.
    k : int
        Number of top hubs to return.

    Returns
    -------
    top_hubs : array of shape (min(k, n),)
        The actual point indices (not positions) of the top k hubs.
    r   r   r   r   r+   minr.   fullr3   r/   r1   )r   rp   kro   actual_ktop_degreestop_indicesr   deg
insert_posrq   s              r?   get_top_k_hub_indicesr~     s,   0 }}QH1hH ''(BHHRLAK((82884K8_ 1WQZ( X\**!AJq.S;zA~+F%Fa
 q.S;zA~+F%F 8a<R8 4!,QU!3A!,QU!3A4
 '*K
#&-ajK
##1& rA   g?c           	         | j                   d   }|j                   d   }t        ||d      }|j                   d   }t        j                  d      }	t        j                  d      }
t        j                  d      }t        j
                  |t        j                        }t        j                  d      }t        j
                  |t        j                        }t        j                  |t        j                        }t        |      D ]  }t        |dz   |      D ]  }||   }||   }t        j                  d      }t        j                  |t        j                        }t        |      D ]/  }| ||f   | ||f   z
  ||<   |||   | ||f   | ||f   z   z  dz  z  }1 t        j                  d      }t        j                  d      }t        |      D ]k  }|}t        |      D ]  }|||   | ||   |f   z  z  } |t        kD  rd||<   |dz  }<|t         k  rd||<   |dz  }Q|dz  ||<   ||   dk(  r|dz  }g|dz  }m |dk(  s|dk(  r4t        j                  t        ||            t        j                  |      z  }||	kD  sp|}	|}
|}|}t        |      D ]
  }||   ||<    t        |      D ]
  }||   ||<      |
dk(  s|dk(  rqt        j                  d      }
t        j                  d      }t        |      D ]9  }t        j                  t        |            dz  ||<   ||   dk(  r|
dz  }
5|dz  }; t        j                  |
t        j                        }t        j                  |t        j                        }t        j                  d      }t        j                  d      }t        |      D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' |||||	fS )a  Hub-based graph-informed split using balance-based selection.

    Uses the top 3 highest-degree nodes to generate all 3 possible hyperplanes,
    then selects the one with the best balance (closest to 50/50 split).
    This is much faster than edge-cut counting while still producing good quality trees.

    Parameters
    ----------
    data : array of shape (n_samples, n_features)
        The data array.
    indices : array of shape (n,)
        Indices of points in this node.
    neighbor_indices : array of shape (n_samples, n_neighbors)
        The neighbor graph.
    global_degrees : array of shape (n_samples,)
        Precomputed global in-degrees.
    rng_state : array of int64, shape (3,)
        RNG state (only used for fallback).

    Returns
    -------
    indices_left, indices_right, hyperplane, offset, balance
        The balance is returned so the caller can decide whether to accept the split.
    r   r      r)   r   rJ   r*   )r+   r~   r.   r0   uint32rm   r2   r/   r1   r-   rv   r,   r	   r3   )r4   r   rn   rp   r5   r6   ro   top_hubsn_hubsbest_balancebest_n_leftbest_n_rightbest_hyperplanebest_offset	best_sider<   ri   hjr7   r8   r   r   r   r   r   r   r   balancer=   r>   s                                 r?   euclidean_hub_splitr   !  s   < **Q-C}}QH %Wna@H^^AF ::c?L))A,K99Q<Lhhs"**5O**S/K1I88HBGG,D Fm 4+Q' 3	+BB<DRLE !#

3 "BJJ ?3Z '+D!G}tE1H~'E!!$!%a(DqMDN,JKcQ! YYq\FiilG8_ %*s IA/2T'!*a-5HHHFI C<DGaKFsd]DGqLG!eDGAw!|!1!%& {gl jjVW!56H9MMG%&$&/s >A):1)=OA&>x +A#'7IaL+e3	+4+n a<1,iilyy|x 	"A66,y"9:Q>IaL|q q !	" 88Krxx8LHH\:MYYq\FiilG8_ Q<1#*1:L aKF%,QZM'"qLG lRRrA   c           	         | j                   d   }|j                   d   }t        ||d      }|j                   d   }t        j                  d      }	t        j                  d      }
t        j                  d      }t        j
                  |t        j                        }t        j
                  |t        j                        }t        j                  |t        j                        }t        |      D ]"  }t        |dz   |      D ]  }||   }||   }t        | |         }t        | |         }t        |      t        k  rd}t        |      t        k  rd}t        j                  |t        j                        }t        |      D ]  }| ||f   |z  | ||f   |z  z
  ||<    t        |      }t        |      t        k  rd}t        |      D ]  }||   |z  ||<    t        j                  d      }t        j                  d      }t        |      D ]~  }t        j                  d      }t        |      D ]  }|||   | ||   |f   z  z  } |t        kD  rd||<   |dz  }O|t         k  rd||<   |dz  }d|dz  ||<   ||   dk(  r|dz  }z|dz  } |dk(  s|dk(  rt        j                  t        ||            t        j                  |      z  }||	kD  s|}	|}
|}t        |      D ]
  }||   ||<    t        |      D ]
  }||   ||<     % |
dk(  s|dk(  rqt        j                  d      }
t        j                  d      }t        |      D ]9  }t        j                  t        |            dz  ||<   ||   dk(  r|
dz  }
5|dz  }; t        j                  |
t        j                        }t        j                  |t        j                        }t        j                  d      }t        j                  d      }t        |      D ]%  }||   dk(  r||   ||<   |dz  }||   ||<   |dz  }' |||t        j                  d      |	fS )a  Angular hub-based split using balance-based selection.

    Uses the top 3 highest-degree nodes to generate all 3 possible hyperplanes,
    then selects the one with the best balance (closest to 50/50 split).

    Returns
    -------
    indices_left, indices_right, hyperplane, offset, balance
        The balance is returned so the caller can decide whether to accept the split.
    r   r   r   r)   r   r(   r*   )r+   r~   r.   r0   r   rm   r2   r/   r1   r
   r,   r-   rv   r	   r3   )r4   r   rn   rp   r5   r6   ro   r   r   r   r   r   r   r   r<   ri   r   r7   r8   r9   r:   r   r   r;   r   r   r   r   r   r=   r>   s                                  r?   angular_hub_splitr     sL     **Q-C}}QH %Wna@H^^AF ::c?L))A,K99Q<Lhhs"**5O1I88HBGG,D Fm >+Q' =	+BB<DRLE T$Z(Id5k*J9~#	:$ 
 "BJJ ?3Z (,T1W	(ANZ/(!!$
 ##45O?#c)"%3Z N'8';o'M!!$N YYq\FiilG8_ %Cs IA/2T'!*a-5HHHFI C<DGaKFsd]DGqLG!eDGAw!|!1!%& {gl jjVW!56H9MMG%&$&s >A):1)=OA&>x +A#'7IaL+y=	+>+B a<1,iilyy|x 	"A66,y"9:Q>IaL|q q !	" 88Krxx8LHH\:MYYq\FiilG8_ Q<1#*1:L aKF%,QZM'"qLG C,VVrA   )left_node_numright_node_num)r$   r%   r"   c                    |j                   d   |	kD  r|
dkD  rt        | ||||      \  }}}}}|t        k  r|j                  t	        j
                  dgt        j                               |j                  t        j                          |j                  t	        j                  d      t	        j                  d      f       |j                  |       yt        | |||||||||	|
dz
         t        |      dz
  }t        | |||||||||	|
dz
         t        |      dz
  }|j                  |       |j                  |       |j                  t	        j                  |      t	        j                  |      f       |j                  t	        j
                  dgt        j                               y|j                  t	        j
                  dgt        j                               |j                  t        j                          |j                  t	        j                  d      t	        j                  d      f       |j                  |       y)zRecursive tree builder using hub-based splits.

    Stops splitting if:
    - Node size <= leaf_size
    - max_depth reached
    - Best split balance < MIN_SPLIT_BALANCE (creates larger leaf instead of bad split)
    r         r   r   Nr   )r+   r   MIN_SPLIT_BALANCEappendr.   arrayr0   infr3   make_hub_euclidean_treelenr4   r   rn   rp   r   r   r   point_indicesr5   r   	max_depthleft_indicesright_indicesr_   offsetr   r   r   s                     r?   r   r     s   2 }}Q)#	A  '+^Y
	
 &&rxxbjjABNNBFF7#OORXXb\288B<89  )M	
 M*Q.M	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288TF"**=>w"rxx|45W%
rA   c                    |j                   d   |	kD  r|
dkD  rt        | ||||      \  }}}}}|t        k  r|j                  t	        j
                  dgt        j                               |j                  t        j                          |j                  t	        j                  d      t	        j                  d      f       |j                  |       yt        | |||||||||	|
dz
         t        |      dz
  }t        | |||||||||	|
dz
         t        |      dz
  }|j                  |       |j                  |       |j                  t	        j                  |      t	        j                  |      f       |j                  t	        j
                  dgt        j                               y|j                  t	        j
                  dgt        j                               |j                  t        j                          |j                  t	        j                  d      t	        j                  d      f       |j                  |       y)zRecursive tree builder using angular hub-based splits.

    Stops splitting if:
    - Node size <= leaf_size
    - max_depth reached
    - Best split balance < MIN_SPLIT_BALANCE (creates larger leaf instead of bad split)
    r   r   r   r   Nr   )r+   r   r   r   r.   r   r0   r   r3   make_hub_angular_treer   r   s                     r?   r   r   w  s   2 }}Q)#	A '+^Y
	
 &&rxxbjjABNNBFF7#OORXXb\288B<89  )M	
 M*Q.M	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288TF"**=>w"rxx|45W%
rA   )r$   r%   c                    t        |      }t        j                  | j                  d         j	                  t        j
                        }t        j                  j                  j                  t              }t        j                  j                  j                  t              }	t        j                  j                  j                  t              }
t        j                  j                  j                  t              }|rt        | |||||	|
||||       nt        | |||||	|
||||       |}|D ]/  }t!        |      |kD  st        j
                  t!        |            }1 t#        ||	|
||      }|S )a  Build an RP tree using simplified hub-based hyperplane selection.

    This version precomputes global degrees once and uses the top 3 highest-degree
    nodes at each split to generate all 3 possible hyperplanes. This is simpler
    and significantly faster than the random sampling approach while maintaining
    or improving tree quality.

    Parameters
    ----------
    data : array of shape (n_samples, n_features)
        The data to build the tree on.
    neighbor_indices : array of shape (n_samples, n_neighbors)
        The neighbor graph indices.
    rng_state : array of int64, shape (3,)
        The internal state of the rng.
    leaf_size : int
        The maximum size of a leaf node.
    angular : bool
        Whether to use angular (cosine) or euclidean distance.
    max_depth : int
        Maximum tree depth.

    Returns
    -------
    tree : FlatTree
        The constructed tree.
    r   r   )rs   r.   aranger+   rP   r3   numbatypedList
empty_listdense_hyperplane_typeoffset_typechildren_typepoint_indices_typer   r   r   r   )r4   rn   r5   r   angularr   rp   r   r   r   r   r   max_leaf_sizepointsresults                  r?   make_hub_treer     sC   J ,,<=Nii

1&--bhh7G++""--.CDKkk))+6G{{**=9HKK$$//0BCM	
 	 	
 M 5v;&!KKF4M5 k7Hm]SFMrA   c           
      ,
   |j                   d   }t        ||d      }|j                   d   }	t        j                  |j                   d   dt        j                        }
t        |      D ]
  }||
||   <    t        j                  t        j                  d      g      }t        j                  t        j                  d      g      }t        j                  d      }t        j                  |t        j                        }t        j                  d      }t        j                  d      }t        j                  d      }t        j                  |t        j                        }t        |	      D ]a  }t        |dz   |	      D ]K  }||   }||   }| ||   ||dz       }|||   ||dz       }| ||   ||dz       }|||   ||dz       }t        j                  d      }t        ||||      \  }}t        ||||      \  }} | d	z  } t        |||| j                  t        j                              \  }} | D ]  }!||!z  }	 t        j                  d      }"t        j                  d      }#t        |      D ]  }|}$| |||      |||   dz       }%||||      |||   dz       }&t        |||%|&      \  }'}(|(D ]  }!|$|!z  }$	 |$t         kD  rd||<   |"dz  }"^|$t          k  rd||<   |#dz  }#s|d
z  ||<   ||   dk(  r|"dz  }"|#dz  }# |"dk(  s|#dk(  rt        j                  d      })t        |      D ]P  }||   }*||   }+t        |j                   d         D ])  },||*|,f   }-|-dk  r 6|
|-   }.|.dk\  s||.   |+k7  s%|)dz  })+ R |)d
z  })|)|k  s|)}|"}|#}|j#                         }|j#                         }|}t        |      D ]
  }||   ||<    N d |dk(  s|dk(  rqt        j                  d      }t        j                  d      }t        |      D ]9  }t        j$                  t'        |            d
z  ||<   ||   dk(  r|dz  }5|dz  }; t        j                  |t        j                        }/t        j                  |t        j                        }0t        j                  d      }"t        j                  d      }#t        |      D ]%  }||   dk(  r||   |/|"<   |"dz  }"||   |0|#<   |#dz  }#' t        j(                  ||f      }1|/|0|1|fS )zSimplified hub-based split for sparse euclidean data.

    Uses the top 3 highest-degree nodes to generate 3 possible hyperplanes,
    then selects the one that minimizes edge cuts.
    r   r   r   r   r   r)       r   rJ   r*   )r+   r~   r.   rw   r3   r1   r   r0   float64r/   r2   r   r   r   r   rP   r-   copyr,   r	   rQ   )2rR   rS   spdatar   rn   rp   r5   ro   r   r   
idx_to_posr   best_hyperplane_indsbest_hyperplane_datar   r   best_edge_cutsr   r   r<   ri   r   r7   r8   rT   rU   rV   rW   r   rX   rY   rb   rc   r^   r   r   r   rZ   r[   r\   r]   	edge_cuts	point_idx
point_sidej_nbrr   neighbor_posr=   r>   r_   s2                                                     r?   sparse_euclidean_hub_splitr   (  s    }}QH %Wna@H^^AF )//2BbhhGJ8_ #!"
71:# 88RXXb\N388RZZ%5$67**S/K1IYYz*N))A,K99Q<L88HBGG,D Fm S+Q' R	+BB<DRLEVD\F4!8,<=Ivd|fTAX.>?IfUmfUQY.?@Juuqy0ABJ !#

3/:9j*0,O_ (29j*($K &+K'1""2::.	($K # )!S(!) YYq\FiilG8_ %*fWQZ06'!*q.3IJwqz 2VGAJN5KL(#_ff8 $ "CcMF" C<DGaKFsd]DGqLG!eDGAw!|!1-%0 {gl 		!I8_ 
+#AJ	!!W
!"2"8"8";< +D/	4@H!|#-h#7L#q(-;%NI+
+ "QI>)!*$&'6';';'=$'6';';'=$/x +A#'7IaL+cR	+S+l a<1,iilyy|x 	"A66,y"9:Q>IaL|q q !	" 88Krxx8LHH\:MYYq\FiilG8_ Q<1#*1:L aKF%,QZM'"qLG 02FGHJ
K??rA   c           	         |j                   d   }t        ||d      }|j                   d   }	t        j                  |j                   d   dt        j                        }
t        |      D ]
  }||
||   <    t        j                  t        j                  d      g      }t        j                  t        j                  d      g      }t        j                  |t        j                        }t        j                  d      }t        j                  d      }t        j                  d      }t        j                  |t        j                        }t        |	      D ]  }t        |dz   |	      D ]  }||   }||   }| ||   ||dz       }|||   ||dz       }| ||   ||dz       }|||   ||dz       }t        |      }t        |      }t        |      t        k  rd}t        |      t        k  rd}||z  j                  t        j                        }||z  j                  t        j                        }t        ||||      \  }} t        |       }!t        |!      t        k  rd}!t        | j                   d         D ]  }"| |"   |!z  | |"<    t        j                  d      }#t        j                  d      }$t        |      D ]  }t        j                   d	      }%| |||      |||   dz       }&||||      |||   dz       }'t#        || |&|'      \  }(})|)D ]  }*|%|*z  }%	 |%t        kD  rd||<   |#dz  }#q|%t         k  rd||<   |$dz  }$|d
z  ||<   ||   dk(  r|#dz  }#|$dz  }$ |#dk(  s|$dk(  rt        j                  d      }+t        |      D ]P  }||   },||   }-t        |j                   d         D ])  }.||,|.f   }/|/dk  r 6|
|/   }0|0dk\  s||0   |-k7  s%|+dz  }++ R |+d
z  }+|+|k  s|+}|#}|$}|j%                         }| j%                         }t        |      D ]
  }||   ||<      |dk(  s|dk(  rqt        j                  d      }t        j                  d      }t        |      D ]9  }t        j                  t'        |            d
z  ||<   ||   dk(  r|dz  }5|dz  }; t        j                  |t        j                        }1t        j                  |t        j                        }2t        j                  d      }#t        j                  d      }$t        |      D ]%  }||   dk(  r||   |1|#<   |#dz  }#||   |2|$<   |$dz  }$' t        j(                  ||f      }3|1|2|3t        j                   d	      fS )zSimplified hub-based split for sparse angular data.

    Uses the top 3 highest-degree nodes to generate 3 possible hyperplanes,
    then selects the one that minimizes edge cuts.
    r   r   r   r   r   r   r   r(   r)   r*   )r+   r~   r.   rw   r3   r1   r   r0   r/   r2   r   r
   r,   r-   rP   r   r   r   r   r	   rQ   )4rR   rS   r   r   rn   rp   r5   ro   r   r   r   r   r   r   r   r   r   r   r<   ri   r   r7   r8   rT   rU   rV   rW   r9   r:   rL   rM   rX   rY   r;   r   r   r   r   rZ   r[   r\   r]   r^   r   r   r   r   rr   r   r=   r>   r_   s4                                                       r?   sparse_angular_hub_splitr     s    }}QH %Wna@H^^AF )//2BbhhGJ8_ #!"
71:# 88RXXb\N388RZZ%5$671IYYz*N))A,K99Q<L88HBGG,D Fm V+Q' U	+BB<DRLEVD\F4!8,<=Ivd|fTAX.>?IfUmfUQY.?@Juuqy0ABJ YIj)J9~#	:$ 
$-	$9#A#A"**#M %/*%<$D$DRZZ$P!/:/=R0,O_ #?3O?#c)"%?0034 J%4Q%7/%I"J YYq\FiilG8_ %CfWQZ06'!*q.3IJwqz 2VGAJN5KL(#_ff8 $ "CcMF" C<DGaKFsd]DGqLG!eDGAw!|!1-%0 {gl 		!I8_ 
+#AJ	!!W
!"2"8"8";< +D/	4@H!|#-h#7L#q(-;%NI+
+ "QI>)!*$&'6';';'=$'6';';'=$x +A#'7IaL+iU	+V+r a<1,iilyy|x 	"A66,y"9:Q>IaL|q q !	" 88Krxx8LHH\:MYYq\FiilG8_ Q<1#*1:L aKF%,QZM'"qLG 02FGHJ
BJJsOCCrA   c                 d   |j                   d   |kD  r|dkD  rt        | ||||||
      \  }}}}t        | |||||||||	|
||dz
         t        |	      dz
  }t        | |||||||||	|
||dz
         t        |	      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |	j	                  t        j                  dgt
        j                               y|j	                  t        j                  dgdggt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |	j	                  |       y)zDRecursive tree builder using simplified sparse euclidean hub splits.r   r   r   r   r   N)
r+   r   make_sparse_hub_euclidean_treer   r   r.   r3   r   r   r   rR   rS   r   r   rn   rp   r   r   r   r   r5   r   r   r   r   r_   r   r   r   s                      r?   r   r   R  s   * }}Q)#	A '&&'+;^Y
	

 	'M	
  M*Q.&M	
  ]+a/:&v-0"((>2JKLRXXrd"((;<  	288dVdV$4BJJGHw"rxx|45W%
rA   c                 d   |j                   d   |kD  r|dkD  rt        | ||||||
      \  }}}}t        | |||||||||	|
||dz
         t        |	      dz
  }t        | |||||||||	|
||dz
         t        |	      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |	j	                  t        j                  dgt
        j                               y|j	                  t        j                  dgdggt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |	j	                  |       y)zBRecursive tree builder using simplified sparse angular hub splits.r   r   r   r   r   N)
r+   r   make_sparse_hub_angular_treer   r   r.   r3   r   r   r   r   s                      r?   r   r     s   * }}Q)#	A %&&'+;^Y
	

 	%M	
  M*Q.$M	
  ]+a/:&v-0"((>2JKLRXXrd"((;<  	288dVdV$4BJJGHw"rxx|45W%
rA   c                    t        |      }t        j                  |j                  d   dz
        j	                  t        j
                        }	t        j                  j                  j                  t              }
t        j                  j                  j                  t              }t        j                  j                  j                  t              }t        j                  j                  j                  t              }|rt        | |||	|||
||||||       nt        | |||	|||
||||||       |}|D ]/  }t!        |      |kD  st        j
                  t!        |            }1 t#        |
||||      }|S )a  Build a sparse RP tree using simplified hub-based hyperplane selection.

    This version precomputes global degrees once and uses the top 3 highest-degree
    nodes at each split to generate all 3 possible hyperplanes.

    Parameters
    ----------
    inds : array
        CSR format index array of the matrix.
    indptr : array
        CSR format index pointer array of the matrix.
    spdata : array
        CSR format data array of the matrix.
    neighbor_indices : array of shape (n_samples, n_neighbors)
        The neighbor graph indices.
    rng_state : array of int64, shape (3,)
        The internal state of the rng.
    leaf_size : int
        The maximum size of a leaf node.
    angular : bool
        Whether to use angular (cosine) or euclidean distance.
    max_depth : int
        Maximum tree depth.

    Returns
    -------
    tree : FlatTree
        The constructed tree.
    r   r   r   )rs   r.   r   r+   rP   r3   r   r   r   r   sparse_hyperplane_typer   r   r   r   r   r   r   )rR   rS   r   rn   r5   r   r   r   rp   r   r   r   r   r   r   r   r   s                    r?   make_sparse_hub_treer     sT   R ,,<=NiiQ!+,33BHH=G++""--.DEKkk))+6G{{**=9HKK$$//0BCM$	
  	'	
  M 5v;&!KKF4M5 k7Hm]SFMrA   r   c                    | j                   d   }t        ||      }t        j                  |t        j                  d      t        j                        }t        j
                  |t        j                        }t        |      D ]y  }|| |      }	|	||dz
     kD  s|dz
  }
|
dkD  r!|	||
dz
     kD  r|
dz  }
|
dkD  r|	||
dz
     kD  rt        |dz
  |
d      D ]  }||dz
     ||<   ||dz
     ||<    |	||
<   | |   ||
<   { |S )zGet the indices of the top k highest-degree points for bit data.

    Also returns the pair with maximum Hamming distance among the top k.
    r   r   r   r   ru   )r   r4   rp   rx   ro   ry   rz   r{   r   r|   r}   rq   s               r?   get_top_k_hub_indices_bitr   R  s*    }}QH1hH ''(BHHRLAK((82884K8_ 1WQZ(X\**!AJq.S;zA~+F%Fa
 q.S;zA~+F%F 8a<R8 4!,QU!3A!,QU!3A4 '*K
#&-ajK
#1 rA   c           
      V	   | j                   d   }|j                   d   }t        || |d      }|j                   d   }t        j                  |j                   d   dt        j                        }	t        |      D ]
  }
|
|	||
   <    t        j                  |dz  t        j                        }t        j                  |t        j                        }t        j                  d      }t        j                  d      }t        j                  d      }t        j                  |t        j                        }t        |      D ])  }t        |dz   |      D ]  }||   }||   }t        j                  |dz  t        j                        }|d| }||d }t        |      D ]+  }| ||f   | ||f   z  }|| ||f   z  ||<   || ||f   z  ||<   - t        j                  d      }t        j                  d      }t        |      D ]  }
t        j                  d	      }t        |      D ]6  }|t        ||   | ||
   |f   z     z  }|t        ||   | ||
   |f   z     z  }8 |t        kD  rd||
<   |dz  }p|t         k  rd||
<   |dz  }|
dz  ||
<   ||
   dk(  r|dz  }|dz  } |dk(  s|dk(  r\t        j                  d      }t        |      D ]P  }
||
   }||
   }t        |j                   d         D ])  } ||| f   }!|!dk  r 6|	|!   }"|"dk\  s||"   |k7  s%|dz  }+ R |dz  }||k  s|}|}|}t        |dz        D ]
  }||   ||<    t        |      D ]
  }
||
   ||
<     , |dk(  s|dk(  rqt        j                  d      }t        j                  d      }t        |      D ]9  }
t        j                  t        |            dz  ||
<   ||
   dk(  r|dz  }5|dz  }; t        j                  |t        j                        }#t        j                  |t        j                        }$t        j                  d      }t        j                  d      }t        |      D ]%  }
||
   dk(  r||
   |#|<   |dz  }||
   |$|<   |dz  }' |#|$|t        j                  d	      fS )
zSimplified hub-based split for bit-packed data.

    Uses the top 3 highest-degree nodes to generate 3 possible hyperplanes,
    then selects the one that minimizes edge cuts.
    r   r   r   r   r   r*   r   Nr)   )r+   r   r.   rw   r3   r1   rm   rC   r/   r2   r   r0   rD   r-   r,   r	   )%r4   r   rn   rp   r5   r6   ro   r   r   r   r   r   r   r   r   r   r<   ri   r   r7   r8   r   rE   rF   r   rG   r   r   r   r   r   r   r   rr   r   r=   r>   s%                                        r?   bit_hub_splitr   u  s    **Q-C}}QH )$JH^^AF )//2BbhhGJ8_ #!"
71:# hhsQwbhh7O1IYYz*N))A,K99Q<L88HBGG,D Fm D+Q' C	+BB<DRLE !#q A,=ds,C),=cd,C)3Z O!$']T%(^;
3=T1W3M-a03=UAX3N-a0O YYq\FiilG8_ %Cs Af5a84
A;NN F f5a84
A;NN F	 C<DGaKFsd]DGqLG!eDGAw!|!1+%. {gl 		!I8_ 
+#AJ	!!W
!"2"8"8";< +D/	4@H!|#-h#7L#q(-;%NI+
+ "QI>)!*$&sQw >A):1)=OA&>x +A#'7IaL+EC	+D+N a<1,iilyy|x 	"A66,y"9:Q>IaL|q q !	" 88Krxx8LHH\:MYYq\FiilG8_ Q<1#*1:L aKF%,QZM'"qLG CHHrA   c                 x   |j                   d   |	kD  r|
dkD  rt        | ||||      \  }}}}t        | |||||||||	|
dz
         t        |      dz
  }t        | |||||||||	|
dz
         t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y|j	                  t        j                  t        j                  d      gt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y)z7Recursive tree builder using simplified bit hub splits.r   r   r   r   N)
r+   r   make_bit_hub_tree_recursiver   r   r.   r3   r   rC   r   )r4   r   rn   rp   r   r   r   r   r5   r   r   r   r   r_   r   r   r   s                    r?   r   r     s   & }}Q)#	A $)9>9U	
 	$M	
 M*Q.#M	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288RXXa[MBCw"rxx|45W%
rA   c                    t        |      }t        j                  | j                  d         j	                  t        j
                        }t        j                  j                  j                  t              }t        j                  j                  j                  t              }t        j                  j                  j                  t              }	t        j                  j                  j                  t              }
t        | ||||||	|
|||       |}|
D ]/  }t        |      |kD  st        j
                  t        |            }1 t!        |||	|
|      }|S )a  Build a bit-packed RP tree using simplified hub-based hyperplane selection.

    This version precomputes global degrees once and uses the top 3 highest-degree
    nodes at each split to generate all 3 possible hyperplanes.

    Parameters
    ----------
    data : array of shape (n_samples, n_features)
        The bit-packed data to build the tree on.
    neighbor_indices : array of shape (n_samples, n_neighbors)
        The neighbor graph indices.
    rng_state : array of int64, shape (3,)
        The internal state of the rng.
    leaf_size : int
        The maximum size of a leaf node.
    max_depth : int
        Maximum tree depth.

    Returns
    -------
    tree : FlatTree
        The constructed tree.
    r   r   )rs   r.   r   r+   rP   r3   r   r   r   r   bit_hyperplane_typer   r   r   r   r   r   )r4   rn   r5   r   r   rp   r   r   r   r   r   r   r   r   s                 r?   make_bit_hub_treer   >  s   @ ,,<=Nii

1&--bhh7G++""--.ABKkk))+6G{{**=9HKK$$//0BCM M 5v;&!KKF4M5 k7Hm]SFMrA   )r$   r"   c	                 F   |j                   d   |kD  r|dkD  rt        | ||      \  }	}
}}t        | |	|||||||dz
  	       t        |      dz
  }t        | |
|||||||dz
  	       t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y |j	                  t        j                  dgt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y Nr   r   r   r   r   )
r+   rK   make_euclidean_treer   r   r.   r3   r   r0   r   r4   r   r   r   r   r   r5   r   r   r   r   r_   r   r   r   s                  r?   r   r   }  sr    }}Q)#	A .dGYG	
 	M
	
 M*Q.M
	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288TF"**=>w"rxx|45W%
rA   )r   r   r   c	                 F   |j                   d   |kD  r|dkD  rt        | ||      \  }	}
}}t        | |	|||||||dz
  	       t        |      dz
  }t        | |
|||||||dz
  	       t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y |j	                  t        j                  dgt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y r   )
r+   r@   make_angular_treer   r   r.   r3   r   r0   r   r   s                  r?   r   r     sr   & }}Q)#	A ,D'9E	
 	M
	
 M*Q.M
	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288TF"**=>w"rxx|45W%
rA   c	                 F   |j                   d   |kD  r|dkD  rt        | ||      \  }	}
}}t        | |	|||||||dz
  	       t        |      dz
  }t        | |
|||||||dz
  	       t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y |j	                  t        j                  dgt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y )Nr   r   r   r      )
r+   rH   make_bit_treer   r   r.   r3   r   rC   r   r   s                  r?   r   r   	  sr   & }}Q)#	A 6dGYO	
 	M
	
 M*Q.M
	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288SE:;w"rxx|45W%
rA   c                 X   |j                   d   |	kD  r|
dkD  rt        | ||||      \  }}}}t        | |||||||||	|
dz
         t        |      dz
  }t        | |||||||||	|
dz
         t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y |j	                  t        j                  dgdggt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y r   )
r+   rd   make_sparse_euclidean_treer   r   r.   r3   r   r   r   rR   rS   r4   r   r   r   r   r   r5   r   r   r   r   r_   r   r   r   s                    r?   r   r   E	  s   " }}Q)#	A 5&$
	

 	#M	
 M*Q."M	
 ]+a/:&v-0"((>2JKLRXXrd"((;<  	288dVdV$4BJJGHw"rxx|45W%
rA   c                 X   |j                   d   |	kD  r|
dkD  rt        | ||||      \  }}}}t        | |||||||||	|
dz
         t        |      dz
  }t        | |||||||||	|
dz
         t        |      dz
  }|j	                  |       |j	                  |       |j	                  t        j                  |      t        j                  |      f       |j	                  t        j                  dgt
        j                               y |j	                  t        j                  dgdggt
        j                               |j	                  t
        j                          |j	                  t        j                  d      t        j                  d      f       |j	                  |       y r   )
r+   r`   make_sparse_angular_treer   r   r.   r3   r   r   r   r   s                    r?   r   r   	  s   " }}Q)#	A 3&$
	

 	!M	
 M*Q. M	
 ]+a/:&v-0"((>2JKLRXXrd"((;<288dVdV$4BJJGHw"rxx|45W%rA   )r$   c                    t        j                  | j                  d         j                  t         j                        }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }	|rt        | |||||	|||	       nt        | |||||	|||	       |}
|	D ]/  }t        |      |
kD  st        j                  t        |            }
1 t!        ||||	|
      }|S )Nr   r   )r.   r   r+   rP   r3   r   r   r   r   r   r   r   r   r   r   r   r   r4   r5   r   r   r   r   r   r   r   r   r   r   r   s                r?   make_dense_treer   	  s)   ii

1&--bhh7G++""--.CDKkk))+6G{{**=9HKK$$//0BCM
	
 	
	
 M 5v;&!KKF4M5 k7Hm]SFMrA   c                    t        j                  |j                  d   dz
        j                  t         j                        }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }	t
        j                  j                  j                  t              }
t
        j                  j                  j                  t              }|rt        | |||||	|
||||       nt        | |||||	|
||||       |}|D ]/  }t        |      |kD  st        j                  t        |            }1 t!        ||	|
||      S )Nr   r   r   )r.   r   r+   rP   r3   r   r   r   r   r   r   r   r   r   r   r   r   )rR   rS   r   r5   r   r   r   r   r   r   r   r   r   r   s                 r?   make_sparse_treer   	  s8    iiQ!+,33BHH=G++""--.DEKkk))+6G{{**=9HKK$$//0BCM 	
 	#	
 M 5v;&!KKF4M5 K(M=QQrA   c                    t        j                  | j                  d         j                  t         j                        }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }t
        j                  j                  j                  t              }	|rt        | |||||	|||	       nt        d      |}
|	D ]/  }t        |      |
kD  st        j                  t        |            }
1 t!        ||||	|
      }|S )Nr   r   z,Euclidean bit trees are not implemented yet.)r.   r   r+   rP   r3   r   r   r   r   r   r   r   r   r   NotImplementedErrorr   r   r   s                r?   make_dense_bit_treer   3
  s   ii

1&--bhh7G++""--.ABKkk))+6G{{**=9HKK$$//0BCM
	
 ""PQQM 5v;&!KKF4M5 k7Hm]SFMrA   zb1(f4[::1],f4,f4[::1],i8[::1])C)readonly)r   r6   r   )r#   r"   r%   c                     |}|j                   d   }t        |      D ]  }|| |   ||   z  z  } t        |      t        k  r(t	        j                  t        |            dz  }|dk(  ryy|dkD  ryyNr   r*   r   )r+   r1   r,   r-   r.   r	   r_   r   pointr5   r   r6   r   r<   s           r?   select_sider   T
  s    & F
++a.C3Z +*Q-%(**+ 6{Svvl9-.219	!rA   zb1(u1[::1],f4,u1[::1],i8[::1])c                     |}|j                   d   }t        |      D ]/  }|t        | |   ||   z     z  }|t        | ||z      ||   z     z  }1 t        |      t        k  r(t        j                  t        |            dz  }|dk(  ryy|dkD  ryyr   )r+   r1   rD   r,   r-   r.   r	   r   s           r?   select_side_bitr   x
  s    & F
++a.C3Z 9&Aq122&C!G,uQx7889 6{Svvl9-.219	!rA   z<i4[::1](f4[::1],f4[:,::1],f4[::1],i4[:,::1],i4[::1],i8[::1])r*   )noder<   )r"   r%   c                     d}||df   dkD  r3t        ||   ||   | |      }|dk(  r||df   }n||df   }||df   dkD  r3|||df    ||df     S Nr   r   )r   r   r   r   r   r   r5   r   r<   s           r?   search_flat_treer   
  s      D
47
a
;t,gdmUIN19D!G$DD!G$D 47
a
 HT1W%%$'):(:;;rA   z<i4[::1](u1[::1],u1[:,::1],f4[::1],i4[:,::1],i4[::1],i8[::1])c                     d}||df   dkD  r3t        ||   ||   | |      }|dk(  r||df   }n||df   }||df   dkD  r3|||df    ||df     S r   )r   r   s           r?   search_flat_bit_treer   
  s      D
47
a
{40'$-	R19D!G$DD!G$D 47
a
 HT1W%%$'):(:;;rA   )r#   r%   c                 @   |}| j                   d   }| d|dz
  f   dk  r|dz  }| d|dz
  f   dk  r| dd |f   j                  t        j                        }| dd |f   }|t	        ||||      z  }t        |      t        k  rt        |      dz  }	|	dk(  ryy|dkD  ryy)Nr   r   r)   r*   )r+   rP   r.   r3   r   r,   r-   r	   )
r_   r   
point_inds
point_datar5   r   hyperplane_sizerX   rY   r<   s
             r?   sparse_select_sider   
  s    F &&q)O
Q!++
,s
21 Q!++
,s
2 !$4_$4!45<<RXXFO $4_$4!45O
 *j F 6{SI&*19	!rA   r   c                     d}||df   dkD  r4t        ||   ||   | ||      }|dk(  r||df   }n||df   }||df   dkD  r4|||df    ||df     S r   )r   )	r   r   r   r   r   r   r5   r   r<   s	            r?   search_sparse_flat_treer   
  s     D
47
a
!wt}j*i
 19D!G$DD!G$D 47
a
 HT1W%%$'):(:;;rA   c
           
         	 g }
,t        dt        ddt        j                  |      z              |d}|j	                  t
        t        |df      j                  t        j                        	 t        j                  j                         r4 t        j                  |d       	fd	t        |      D              }
ni|r4 t        j                  |d       	fd
t        |      D              }
n3 t        j                  |d       	fdt        |      D              }
t'        |
      S # t        t         t"        f$ r t%        d       Y t'        |
      S w xY w)zBuild a random projection forest with ``n_trees``.

    Parameters
    ----------
    data
    n_neighbors
    n_trees
    leaf_size
    rng_state
    angular

    Returns
    -------
    forest: list
        A list of random projection trees.
    <   r      r   r   )size	sharedmemn_jobsrequirec           
   3      K   | ]K  } t        j                  t              j                  j                  j
                  |           M ywr   N)joblibdelayedr   r   rS   r4   .0r   r   r4   r   r   
rng_statess     r?   	<genexpr>zmake_forest.<locals>.<genexpr>'  sX      I  1/0LLKKIIqM' Is   AAc              3   l   K   | ]+  } t        j                  t              |           - ywr	  )r
  r  r   r  s     r?   r  zmake_forest.<locals>.<genexpr>4  sB      I  423*Q-Gy I   14c              3   l   K   | ]+  } t        j                  t              |           - ywr	  )r
  r  r   r  s     r?   r  zmake_forest.<locals>.<genexpr>;  sA      I  0/*Q-Gy Ir  zRandom Projection forest initialisation failed due to recursionlimit being reached. Something is a little strange with your graph_data, and this may take longer than normal to compute.)maxrv   r.   r3   randint	INT32_MIN	INT32_MAXrP   int64scipysparseisspmatrix_csrr
  Parallelr1   RuntimeErrorRecursionErrorSystemErrorr   tuple)r4   n_neighborsn_treesr   r5   random_stater  r   bit_treer   r   r  s   `  `   ` ` @r?   make_forestr$  
  sI   : FCQ+)>%>?@	~%%i'1%NUU
J!
<<&&t,HV__FKH I wI F HV__FKH I w	I F IV__FKH I w	I F = .+6 
K	
 =
s   6B<D= =E)(E)c                    d}t        t        | j                              D ]3  }| j                  |   d   dk(  s| j                  |   d   dk(  s/|dz  }5 t        j                  ||fdt        j
                        }d}t        t        | j                              D ]d  }| j                  |   d   dk(  s| j                  |   d   dk(  s.| j                  |   j                  d   }| j                  |   ||d |f<   |dz  }f |S )Nr   r   r   r   )r1   r   r   r.   rw   r3   r   r+   )treer   n_leavesr   r   
leaf_indexr   s          r?   get_leaves_from_treer)  K  s   H3t}}%& ==A"$q)9!)<)BMH WWh."((CFJ3t||$% ==A"$a(8(;r(AQ--a0I-1\\!_F:z	z)*!OJ	 MrA   c                     t        j                  | D cg c]  }|j                   c}       t        j                  dd      fd| D              }|S c c}w )Nr   r  r  c              3   ^   K   | ]$  } t        j                  t              |       & y wN)r
  r  r)  )r  rp_treer   s     r?   r  z-rptree_leaf_array_parallel.<locals>.<genexpr>_  s,      = 	-+,WmD=s   *-)r.   r  r   r
  r  )	rp_forestr-  r   r   s      @r?   rptree_leaf_array_parallelr/  ]  sS    FFYG'G--GHM<V__B< = = F M Hs   Ac                     t        |       dkD  rt        j                  t        |             S t        j                  dgg      S )Nr   r   )r   r.   rQ   r/  r   )r.  s    r?   rptree_leaf_arrayr1  f  s6    
9~yy3I>??xx"rA   c                    | j                   |   d   dk  rA|t        | j                  |         z   }| ||df<   | ||df<   | j                  |   ||| ||fS | j                  |   ||<   | j                  |   ||<   |dz   ||df<   |}	t        | |||||dz   || j                   |   d         \  }}|dz   ||	df<   t        | |||||dz   || j                   |   d         \  }}||fS r   )r   r   r   r   r   recursive_convert
r&  r   r   r   r   node_num
leaf_start	tree_nodeleaf_endold_node_nums
             r?   r3  r3  n  sD    }}Y"Q&DLL$; <<!+1!)	1'+||I'>
8$!! $ 0 0 ;H LL3 (110qLMM)$Q'	 
* %-qLq!0qLMM)$Q'	 
* ##rA   c                    | j                   |   d   dk  rA|t        | j                  |         z   }| ||df<   | ||df<   | j                  |   ||| ||fS | j                  |   ||d d d | j                  |   j                  d   f<   | j
                  |   ||<   |dz   ||df<   |}	t        | |||||dz   || j                   |   d         \  }}|dz   ||	df<   t        | |||||dz   || j                   |   d         \  }}||fS r   )r   r   r   r   r+   r   recursive_convert_sparser4  s
             r?   r;  r;    sn    }}Y"Q&DLL$; <<!+1!)	1'+||I'>
8$!! Y' 	Ha!G4#3#3I#>#D#DQ#G!GGH !LL3 (117qLMM)$Q'	 
* %-qLq!7qLMM)$Q'	 
* ##rA   )r%   c                     d}d}t        t        | j                              D ]'  }| j                  |   d   dk  r|dz  }|dz  }#|dz  }) ||fS r   )r1   r   r   )r&  n_nodesr'  r   s       r?   num_nodes_and_leavesr>    sg    GH3t}}%& ==A"MHqLGqLG HrA   c                    t        |       \  }}d}| j                  d   j                  dk(  rc| j                  d   j                  t        j
                  k(  r|dz  }n|}t	        j                  ||f| j                  d   j                        }n8d}|}t	        j                  |d|ft        j                        }d|d d dd d f<   t	        j                  |t        j                        }t	        j                  d      t	        j                  |dft        j                        z  }	t	        j                  d      t	        j                  |t        j                        z  }
|r)t        | |||	|
ddt        | j                        dz
         n(t        | |||	|
ddt        | j                        dz
         t        |||	|
| j                        S )NFr   r   r*   r   Tr   )r>  r   ndimr   r.   rC   rm   r0   r3   onesr;  r   r   r3  r   r   )r&  	data_sizedata_dimr=  r'  	is_sparsehyperplane_dimr   r   r   r   s              r?   convert_tree_formatrF    s   ,T2GXI1$A$$0%\N%Nhhn%T-=-=a-@-F-F

 	!hhN;2::N!Aq!Ghhwbjj1Gxx|bggwl"((CCHhhrlRWWYbhh??G +w'1aT]]ASVWAW	
 	+w'1aT]]ASVWAW	
 K(GT^^LLrA      c                 x    | j                   | j                  | j                  | j                  | j                  f}|S r,  r   r&  r   s     r?   denumbaify_treerJ    s5    F MrA   c                 j    t        | t           | t           | t           | t           | t
                 }|S r,  )r   FLAT_TREE_HYPERPLANESFLAT_TREE_OFFSETSFLAT_TREE_CHILDRENFLAT_TREE_INDICESFLAT_TREE_LEAF_SIZErI  s     r?   renumbaify_treerQ     s@    "#  !F MrA   )intersectionr   r   )parallelr"   r%   c           	         d}t        j                  |j                  d         D ]t  }t        ||   | j                  | j
                  | j                  | j                  |      }t        ||   |      }|t        j                  |j                  d   dkD        z  }v |t        j                  |j                  d         z  S )Nr)   r   r   )
r   pranger+   r   r   r   r   r   r   r0   )r&  rn   r4   r5   r   r   leaf_indicesrR  s           r?   
score_treerW    s     F\\*0034 
;'GLLMMLL
 %%5a%8,G%-- 2 21 5 9::
; EMM"2"8"8";<<<rA   )r   count)r$   r"   r%   c                 ~   |j                   d   }|j                   d   }d}t        | j                        }t        |      D ]  }t	        j
                  |      }| j                  |   d   }| j                  |   d   }	|dk(  sB|	dk(  sH| j                  |   }
|
j                   d   }t        |      D ]p  }|
|   }||   }d}t        |      D ]&  }||   }t        |      D ]  }|
|   |k(  s|dz  } & ( |t	        j                  |      t	        j                  |      z  z  }r  |t	        j                  |      z  S )a~  Score a tree by measuring how well leaves contain nearest neighbors.

    For each point, computes the fraction of its k nearest neighbors that
    are in the same leaf. Returns the average of this fraction across all points.

    A score of 1.0 means all neighbors are always in the same leaf (perfect).
    A score of 0.0 means no neighbors are ever in the same leaf (worst).
    r   r   r)   r   )r+   r   r   r1   r   r3   r   r0   )r&  rn   ro   rx   total_scorer=  r   r   
left_childright_childrV  r   rq   idx	neighborsrX  nirr   lis                      r?   score_linked_treera  %  sc     %%a(Hq!AK$-- G7^ G{{1~]]4(+
mmD)!, r 1<<-L$**1-I 9% G"1o,S1	 ( "B(}H#I. "'+x7!QJE!"" u}}U3emmA6FFF%GG@ x000rA   )r  )      )rb  Frc  )r   )NFFrc  )kwarningsr   numpyr.   r   scipy.sparser  pynndescent.sparser   r   r   r   r   pynndescent.utilsr	   r
   r
  collectionsr   r-   iinfor3   rv   r  r  r  r   r0   r   r   r   rC   r   r   typeofr   r   r   r1   binrX  rD   njittypesTupler  r   r@   rH   rK   r`   rd   FAST_SPLIT_THRESHOLDrk   rs   r~   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   ListTyper   r   r   r   r   r   r   booleanArrayintpuint16r   r   r   r   r   r   r$  r)  r/  r1  r3  r;  r>  rF  rL  rM  rN  rO  rP  rJ  rQ  rW  ra  )r   s   0r?   <module>rv     s9        1  " BHHRXX""Q&	BHHRXX""Q&	N cc* q#A#v. kk#A#& mmhbhhrlHBHHRL9:[[1% 	eCj93q6<<$9	L EKK	SqS	5;;ss+-BKPmmAssFU[[1-u{{3Q3/?A ,,<<"]]3Q3/"]]--\\\\ll||
 

#&o?'&o?d EKK	SqS	5;;ss+-@+Nkk!SqS&5;;ss+U[[1-=? ,,<<"[[1-"]]--\\\\ll||
 

#&l?'&l?^ EKK	SqS	5;;ss+-BKPmmAssFU[[1-u{{3Q3/?A ,,<<"]]3Q3/"]]--\\\\ll||
 

#&aM'&aMH 

 % 3 3CaC 8!&!4!4SqS!9 ;;..[[		
|8
|8~ TT2sF 3sFx   



 



8 


.
.f   


yS
ySx 


tW
tWn 

"[[..%++BSBST Q
Qh 

"[[..%++BSBST Q
Qh $e$
 O %On 


L@
L@^ 


ND
NDb 

"[[..%++BSBST" H
HV 

"[[..%++BSBST" H
HV $e$ W %W~ 



< 


yI
yIx 

"[[..%++BSBST @
@F $e$
 ; %;| 
"[[..%++BSBST 9	9x 
KK((7**++++  99x 
KK((7**++++  99x 
"[[..%++BSBST A	AH 
"[[..%++BSBST ?&	?&D $& &R $ 2R 2Rj $ @ (KKekk111cDIKKKKekk111cDIKKekk//C%H		
 ++%%{{[[
 !$%$$ (KKekk//C$GKKKKekk//C$GKKekk//C%H		
 ++%%{{[[
 !$%$& FC%++++QdCKKekk111cDIKKekk111cDIKKekk111cDIKKekk//C$GKKekk//C$GKKekk//C%H	

 KK&&0C0CD
	<	< FC%++++QdCKKekk//C$GKKekk//C$GKKekk111cDIKKekk//C$GKKekk//C$GKKekk//C%H	

 KK&&0C0CD
	<	< T& '4 FEKK../t<< =<. IX $ " #$L %$ %$P $
 
MD      		 CaC(--\\
 ==  
KK%++6

.1
.1I` :s   4AA,