fn_randperm.hpp 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. // Copyright 2008-2016 Conrad Sanderson (http://conradsanderson.id.au)
  2. // Copyright 2008-2016 National ICT Australia (NICTA)
  3. //
  4. // Licensed under the Apache License, Version 2.0 (the "License");
  5. // you may not use this file except in compliance with the License.
  6. // You may obtain a copy of the License at
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. // ------------------------------------------------------------------------
  15. //! \addtogroup fn_randperm
  16. //! @{
  17. template<typename obj_type>
  18. inline
  19. void
  20. internal_randperm_helper(obj_type& x, const uword N, const uword N_keep)
  21. {
  22. arma_extra_debug_sigprint();
  23. typedef typename obj_type::elem_type eT;
  24. // see op_sort_index_bones.hpp for the definition of arma_sort_index_packet
  25. typedef arma_sort_index_packet<int> packet;
  26. std::vector<packet> packet_vec(N);
  27. for(uword i=0; i < N; ++i)
  28. {
  29. packet_vec[i].val = int(arma_rng::randi<int>());
  30. packet_vec[i].index = i;
  31. }
  32. arma_sort_index_helper_ascend<int> comparator;
  33. if(N >= 2)
  34. {
  35. if(N_keep < N)
  36. {
  37. typename std::vector<packet>::iterator first = packet_vec.begin();
  38. typename std::vector<packet>::iterator nth = first + N_keep;
  39. typename std::vector<packet>::iterator pastlast = packet_vec.end();
  40. std::partial_sort(first, nth, pastlast, comparator);
  41. }
  42. else
  43. {
  44. std::sort( packet_vec.begin(), packet_vec.end(), comparator );
  45. }
  46. }
  47. if(is_Row<obj_type>::value)
  48. {
  49. x.set_size(1,N_keep);
  50. }
  51. else
  52. {
  53. x.set_size(N_keep,1);
  54. }
  55. eT* x_mem = x.memptr();
  56. for(uword i=0; i < N_keep; ++i)
  57. {
  58. x_mem[i] = eT( packet_vec[i].index );
  59. }
  60. }
  61. template<typename obj_type>
  62. arma_warn_unused
  63. inline
  64. typename enable_if2< is_Mat<obj_type>::value, obj_type >::result
  65. randperm(const uword N)
  66. {
  67. arma_extra_debug_sigprint();
  68. obj_type x;
  69. if(N > 0) { internal_randperm_helper(x, N, N); }
  70. return x;
  71. }
  72. arma_warn_unused
  73. inline
  74. uvec
  75. randperm(const uword N)
  76. {
  77. arma_extra_debug_sigprint();
  78. uvec x;
  79. if(N > 0) { internal_randperm_helper(x, N, N); }
  80. return x;
  81. }
  82. template<typename obj_type>
  83. arma_warn_unused
  84. inline
  85. typename enable_if2< is_Mat<obj_type>::value, obj_type >::result
  86. randperm(const uword N, const uword M)
  87. {
  88. arma_extra_debug_sigprint();
  89. arma_debug_check( (M > N), "randperm(): 'M' must be less than or equal to 'N'");
  90. obj_type x;
  91. if( (N > 0) && (M > 0) ) { internal_randperm_helper(x, N, M); }
  92. return x;
  93. }
  94. arma_warn_unused
  95. inline
  96. uvec
  97. randperm(const uword N, const uword M)
  98. {
  99. arma_extra_debug_sigprint();
  100. arma_debug_check( (M > N), "randperm(): 'M' must be less than or equal to 'N'");
  101. uvec x;
  102. if( (N > 0) && (M > 0) ) { internal_randperm_helper(x, N, M); }
  103. return x;
  104. }
  105. //! @}