140 lines
3.8 KiB
C++
140 lines
3.8 KiB
C++
//========= Copyright © Valve Corporation, All rights reserved. ============//
|
|
|
|
|
|
#ifndef MATHLIB_DISJOINT_SET_FOREST_HDR
|
|
#define MATHLIB_DISJOINT_SET_FOREST_HDR
|
|
|
|
#include "tier1/utlvector.h"
|
|
|
|
/// An excellent overview of the concept is here:
|
|
/// http://en.wikipedia.org/wiki/Disjoint-set_data_structure this algorithm is with path
|
|
/// compression and ranking implemented, so it's essentially amortized const-time operations to
|
|
/// find node's island representative element or union two lists. ( the "essentially" means
|
|
/// amortized complexity is Ackermann function, which is like 5 for the largest number in any kind
|
|
/// of software development )
|
|
|
|
class CDisjointSetForest
|
|
{
|
|
public:
|
|
CDisjointSetForest( int nCount );
|
|
|
|
//void Flatten();
|
|
int Find( int nNode );
|
|
void Union( int nNodeA, int nNodeB );
|
|
void EnsureExists( int nNode );
|
|
int GetNodeCount()const { return m_Nodes.Count(); }
|
|
protected:
|
|
struct Node_t
|
|
{
|
|
int nRank, nParent;
|
|
};
|
|
CUtlVector< Node_t > m_Nodes;
|
|
};
|
|
|
|
|
|
inline CDisjointSetForest::CDisjointSetForest( int nCount )
|
|
{
|
|
m_Nodes.SetCount( nCount );
|
|
for( int i = 0;i < nCount; ++i )
|
|
{
|
|
m_Nodes[i].nRank = 0;
|
|
m_Nodes[i].nParent = i;
|
|
}
|
|
}
|
|
|
|
|
|
inline void CDisjointSetForest::EnsureExists( int nNode )
|
|
{
|
|
int nOldCount = m_Nodes.Count();
|
|
if ( nNode >= nOldCount )
|
|
{
|
|
m_Nodes.SetCountNonDestructively( nNode + 1 );
|
|
for ( int n = nOldCount; n <= nNode; ++n )
|
|
{
|
|
m_Nodes[ n ].nRank = 0;
|
|
m_Nodes[ n ].nParent = n;
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
/// Find the representative element for the node in graph representative element is the same for
|
|
/// all connected nodes(vertices) in the graph, and it's one of the nodes in the connected set this
|
|
/// implementation is without recursion to be more cache friendly; recursive implementation would
|
|
/// be clearer, but this is simple enough
|
|
inline int CDisjointSetForest::Find( int nStartNode )
|
|
{
|
|
int nTopParent;
|
|
for( int nNode = nStartNode; nTopParent = m_Nodes[nNode].nParent, nNode != nTopParent ; )
|
|
{
|
|
nNode = nTopParent;
|
|
}
|
|
|
|
// found the top parent, now compress the path to achieve that amazing amortized acceleration
|
|
int nParent;
|
|
for( int nNode = nStartNode; nParent = m_Nodes[nNode].nParent, nNode != nParent ; )
|
|
{
|
|
m_Nodes[nNode].nParent = nTopParent;
|
|
nNode = nParent;
|
|
}
|
|
Assert( nParent == nTopParent );
|
|
return nTopParent;
|
|
}
|
|
|
|
|
|
/// Connect the two (potentially disjoint) sets
|
|
inline void CDisjointSetForest::Union( int nNodeA, int nNodeB )
|
|
{
|
|
int nRootA = Find( nNodeA );
|
|
int nRootB = Find( nNodeB );
|
|
if ( m_Nodes[nRootA].nRank > m_Nodes[nRootB].nRank )
|
|
{
|
|
m_Nodes[nRootB].nParent = nRootA; // note: no change in rank! we're balanced!
|
|
}
|
|
else
|
|
if ( m_Nodes[nRootA].nRank < m_Nodes[nRootB].nRank )
|
|
{
|
|
m_Nodes[nRootA].nParent = nRootB; // note: no change in rank! we're balanced!
|
|
}
|
|
else
|
|
if ( nRootA != nRootB ) // Unless A and B are already in same set, merge them
|
|
{
|
|
m_Nodes[nRootB].nParent = nRootA;
|
|
m_Nodes[nRootA].nRank = m_Nodes[nRootA].nRank + 1;
|
|
}
|
|
}
|
|
|
|
|
|
/// Given the graph implementing GetParent(), find the indices of all children of the given tip of
|
|
/// the subtree
|
|
template <typename Graph_t, class BitVec_t>
|
|
inline void ComputeSubtree( const Graph_t *pGraph, int nSubtreeTipBone, BitVec_t *pSubtree )
|
|
{
|
|
int nBoneCount = pSubtree->GetNumBits();
|
|
Assert( nSubtreeTipBone >= 0 && nSubtreeTipBone < nBoneCount );
|
|
CDisjointSetForest find( nBoneCount );
|
|
for( int nBone = 0; nBone < nBoneCount; ++nBone )
|
|
{
|
|
if( nBone != nSubtreeTipBone ) // Important: severe the link between the subtree tip bone and the rest of the tree to find the disjoint subtree
|
|
{
|
|
int nParent = pGraph->GetParent( nBone );
|
|
if( nParent >= 0 && nParent < nBoneCount )
|
|
{
|
|
find.Union( nBone, nParent );
|
|
}
|
|
}
|
|
}
|
|
int nIsland = find.Find( nSubtreeTipBone );
|
|
for( int nBone = 0; nBone < nBoneCount; ++nBone )
|
|
{
|
|
if( find.Find( nBone ) == nIsland )
|
|
{
|
|
pSubtree->Set( nBone );
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
#endif //MATHLIB_DISJOINT_SET_FOREST_HDR
|
|
|