github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/util/data/ntree.go (about)

     1  package data
     2  
     3  import "sync"
     4  
     5  type NTree struct {
     6  	value    interface{}
     7  	children []*NTree
     8  	mx       sync.Mutex
     9  }
    10  
    11  func (t *NTree) GetValue() interface{} {
    12  	return t.value
    13  }
    14  
    15  func (t *NTree) Find(val interface{}) *NTree {
    16  	t.mx.Lock()
    17  	defer t.mx.Unlock()
    18  
    19  	return t.findUnsafe(val)
    20  }
    21  
    22  func (t *NTree) findUnsafe(val interface{}) *NTree {
    23  	if t.value == val {
    24  		return t
    25  	}
    26  	for _, chld := range t.children {
    27  		if res := chld.Find(val); res != nil {
    28  			return res
    29  		}
    30  	}
    31  	return nil
    32  }
    33  
    34  func (t *NTree) FindParent(val interface{}) *NTree {
    35  	t.mx.Lock()
    36  	defer t.mx.Unlock()
    37  
    38  	return t.findParentUnsafe(val)
    39  }
    40  
    41  func (t *NTree) findParentUnsafe(val interface{}) *NTree {
    42  	for _, child := range t.children {
    43  		if child.value == val {
    44  			return t
    45  		}
    46  	}
    47  	for _, chld := range t.children {
    48  		if res := chld.FindParent(val); res != nil {
    49  			return res
    50  		}
    51  	}
    52  	return nil
    53  }
    54  
    55  func (t *NTree) FindCommonParent(vals ...interface{}) *NTree {
    56  	t.mx.Lock()
    57  	defer t.mx.Unlock()
    58  
    59  	nodes := make([]*NTree, len(vals))
    60  	for ix, val := range vals {
    61  		nodes[ix] = t.findUnsafe(val)
    62  		// If one of the nodes is the root, we return it right away
    63  		if nodes[ix] == t {
    64  			return t
    65  		}
    66  	}
    67  
    68  	lca := nodes[0]
    69  	for ix := 1; ix < len(nodes); ix++ {
    70  		lca = t.findLCA(lca, nodes[ix])
    71  		if lca == t {
    72  			return t
    73  		}
    74  	}
    75  
    76  	return lca
    77  }
    78  
    79  func (t *NTree) findLCA(t1 *NTree, t2 *NTree) *NTree {
    80  	if t == t1 || t == t2 {
    81  		return t
    82  	}
    83  
    84  	finds := make([]*NTree, 0)
    85  	for _, chld := range t.children {
    86  		if res := chld.findLCA(t1, t2); res != nil {
    87  			finds = append(finds, res)
    88  		}
    89  	}
    90  	if len(finds) == 0 {
    91  		return nil
    92  	} else if len(finds) == 1 {
    93  		return finds[0]
    94  	} else if len(finds) == 2 {
    95  		return t
    96  	} else {
    97  		panic("Wrong tree structure: values are found in more than 2 children")
    98  	}
    99  }
   100  
   101  func (t *NTree) FindOrInsert(val interface{}) *NTree {
   102  	t.mx.Lock()
   103  	defer t.mx.Unlock()
   104  
   105  	if ptr := t.findUnsafe(val); ptr == nil {
   106  		ptr = &NTree{value: val}
   107  		t.children = append(t.children, ptr)
   108  		return ptr
   109  	} else {
   110  		return ptr
   111  	}
   112  }
   113  
   114  func (t *NTree) Detach(val interface{}) *NTree {
   115  	t.mx.Lock()
   116  	defer t.mx.Unlock()
   117  
   118  	if ptr := t.findParentUnsafe(val); ptr != nil {
   119  		for ix, chld := range ptr.children {
   120  			if val == chld.value {
   121  				res := ptr.children[ix]
   122  				ptr.children = append(ptr.children[:ix], ptr.children[ix+1:]...)
   123  				return res
   124  			}
   125  		}
   126  	}
   127  	return nil
   128  }
   129  
   130  func (t *NTree) PostTraversal() []interface{} {
   131  	t.mx.Lock()
   132  	defer t.mx.Unlock()
   133  
   134  	stack := make([]interface{}, 0)
   135  	for _, chld := range t.children {
   136  		stack = append(stack, chld.PostTraversal()...)
   137  	}
   138  	if t.value != nil {
   139  		stack = append(stack, t.value)
   140  	}
   141  	return stack
   142  }
   143  
   144  func (t *NTree) PreTraversal() []interface{} {
   145  	t.mx.Lock()
   146  	defer t.mx.Unlock()
   147  
   148  	stack := make([]interface{}, 0)
   149  	if t.value != nil {
   150  		stack = append(stack, t.value)
   151  	}
   152  	for _, chld := range t.children {
   153  		stack = append(stack, chld.PreTraversal()...)
   154  	}
   155  	return stack
   156  }