random.h 2.69 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
2
3
4
5
6
7
#ifndef LIGHTGBM_UTILS_RANDOM_H_
#define LIGHTGBM_UTILS_RANDOM_H_

#include <cstdint>

#include <random>
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
8
#include <set>
Guolin Ke's avatar
Guolin Ke committed
9
10
11
12
13
14
15
16
17
18
19

namespace LightGBM {

/*!
* \brief A wrapper for random generator
*/
class Random {
public:
  /*!
  * \brief Constructor, with random seed
  */
Guolin Ke's avatar
Guolin Ke committed
20
  Random() {
Guolin Ke's avatar
Guolin Ke committed
21
    std::random_device rd;
Guolin Ke's avatar
Guolin Ke committed
22
23
    auto genrator = std::mt19937(rd());
    std::uniform_int_distribution<int> distribution(0, x);
Guolin Ke's avatar
Guolin Ke committed
24
    x = distribution(genrator);
Guolin Ke's avatar
Guolin Ke committed
25
26
27
28
  }
  /*!
  * \brief Constructor, with specific seed
  */
Guolin Ke's avatar
Guolin Ke committed
29
  Random(int seed) {
Guolin Ke's avatar
Guolin Ke committed
30
    x = seed;
Guolin Ke's avatar
Guolin Ke committed
31
32
  }
  /*!
Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
37
38
39
40
41
42
43
  * \brief Generate random integer, int16 range. [0, 65536]
  * \param lower_bound lower bound
  * \param upper_bound upper bound
  * \return The random integer between [lower_bound, upper_bound)
  */
  inline int NextShort(int lower_bound, int upper_bound) {
    return (RandInt16()) % (upper_bound - lower_bound) + lower_bound;
  }

  /*!
  * \brief Generate random integer, int32 range
Guolin Ke's avatar
Guolin Ke committed
44
45
46
47
  * \param lower_bound lower bound
  * \param upper_bound upper bound
  * \return The random integer between [lower_bound, upper_bound)
  */
Guolin Ke's avatar
Guolin Ke committed
48
  inline int NextInt(int lower_bound, int upper_bound) {
Guolin Ke's avatar
Guolin Ke committed
49
    return (RandInt32()) % (upper_bound - lower_bound) + lower_bound;
Guolin Ke's avatar
Guolin Ke committed
50
  }
Guolin Ke's avatar
Guolin Ke committed
51

Guolin Ke's avatar
Guolin Ke committed
52
53
54
55
  /*!
  * \brief Generate random float data
  * \return The random float between [0.0, 1.0)
  */
Guolin Ke's avatar
Guolin Ke committed
56
  inline float NextFloat() {
Guolin Ke's avatar
Guolin Ke committed
57
    // get random float in [0,1)
Guolin Ke's avatar
Guolin Ke committed
58
    return static_cast<float>(RandInt16()) / (32768.0f);
Guolin Ke's avatar
Guolin Ke committed
59
60
61
62
63
64
65
  }
  /*!
  * \brief Sample K data from {0,1,...,N-1}
  * \param N
  * \param K
  * \return K Ordered sampled data from {0,1,...,N-1}
  */
66
67
  inline std::vector<int> Sample(int N, int K) {
    std::vector<int> ret;
68
    ret.reserve(K);
Guolin Ke's avatar
Guolin Ke committed
69
    if (K > N || K <= 0) {
Guolin Ke's avatar
Guolin Ke committed
70
      return ret;
71
72
    } else if (K == N) {
      for (int i = 0; i < N; ++i) {
Guolin Ke's avatar
Guolin Ke committed
73
74
        ret.push_back(i);
      }
Guolin Ke's avatar
Guolin Ke committed
75
    } else if (K > 1 && K > (N / std::log2(K))) {
76
77
78
79
80
81
82
      for (int i = 0; i < N; ++i) {
        double prob = (K - ret.size()) / static_cast<double>(N - i);
        if (NextFloat() < prob) {
          ret.push_back(i);
        }
      }
    } else {
Guolin Ke's avatar
Guolin Ke committed
83
      std::set<int> sample_set;
84
      while (static_cast<int>(sample_set.size()) < K) {
Guolin Ke's avatar
Guolin Ke committed
85
86
87
88
89
90
91
        int next = RandInt32() % N;
        if (sample_set.count(next) == 0) {
          sample_set.insert(next);
        }
      }
      for (auto iter = sample_set.begin(); iter != sample_set.end(); ++iter) {
        ret.push_back(*iter);
92
      }
Guolin Ke's avatar
Guolin Ke committed
93
94
95
    }
    return ret;
  }
96

Guolin Ke's avatar
Guolin Ke committed
97
private:
Guolin Ke's avatar
Guolin Ke committed
98
  inline int RandInt16() {
Guolin Ke's avatar
Guolin Ke committed
99
    x = (214013 * x + 2531011);
Guolin Ke's avatar
Guolin Ke committed
100
    return static_cast<int>((x >> 16) & 0x7FFF);
Guolin Ke's avatar
Guolin Ke committed
101
  }
Guolin Ke's avatar
Guolin Ke committed
102
103
104

  inline int RandInt32() {
    x = (214013 * x + 2531011);
Guolin Ke's avatar
Guolin Ke committed
105
    return static_cast<int>(x & 0x7FFFFFFF);
Guolin Ke's avatar
Guolin Ke committed
106
107
  }

108
  unsigned int x = 123456789;
Guolin Ke's avatar
Guolin Ke committed
109
110
};

Guolin Ke's avatar
Guolin Ke committed
111

Guolin Ke's avatar
Guolin Ke committed
112
113
}  // namespace LightGBM

Guolin Ke's avatar
Guolin Ke committed
114
#endif   // LightGBM_UTILS_RANDOM_H_