from ROOT.VecOps import RVec, Take, Combinations
# RVec can be sorted in Python with the inbuilt sorting function because
# PyROOT implements a Python iterator
v1 = RVec("double")(3)
v1[0], v1[1], v1[2] = 1, 2, 3
v2 = RVec("double")(2)
v2[0], v2[1] = -4, -5
# To get the indices, which result in all combinations, you can call the
# following helper.
# Note that you can also pass the size of the vectors directly.
idx = Combinations(v1, v2)
# Next, the respective elements can be taken via the computed indices.
c1 = Take(v1, idx[0])
c2 = Take(v2, idx[1])
# Finally, you can perform any set of operations conveniently.
v3 = c1 * c2
print("Combinations of {} and {}:".format(v1, v2))
for i in range(len(v3)):
print("{} * {} = {}".format(c1[i], c2[i], v3[i]))
print
# However, if you want to compute operations on unique combinations of a
# single RVec, you can perform this as follows.
# Get the indices of unique triples for the given vector.
v4 = RVec("double")(4)
# an RVec ...
v2 = Sort(v1)
print("Sort vector {}: {}".format(v1, v2))
# ... or a reversed copy of an RVec.
v2 = Reverse(v1)
print("Reverse vector {}: {}".format(v1, v2))
# Helpers are provided to get the indices that sort the vector and to
# select these indices from an RVec.
v2 = Argsort(v1)
print("Indices that sort the vector {}: {}".format(v1, v2))
v3 = RVec("double")(3)
v3[0], v3[1], v3[2] = 9, 7, 8
v4 = Take(v3, v2)
print("Sort vector {} respective to the previously determined indices: {}".
format(v3, v4))
# Take can also be used to get the first or last elements of an RVec.
v2 = Take(v1, 2)
v3 = Take(v1, -2)
print("Take the two first and last elements of vector {}: {}, {}".format(
v1, v2, v3))
# Because the VecOps helpers return a copy of the input, you can chain the operations
# conveniently.
v2 = Reverse(Take(Sort(v1), -2))
print(
"Sort the vector {}, take the two last elements and reverse the selection: {}"
.format(v1, v2))