summaryrefslogtreecommitdiff
path: root/gi/clda/src/crp.h
blob: 9d35857eb5e9fdadc3282f8ac15bb9c80a8f74ba (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#ifndef _CRP_H_
#define _CRP_H_

// shamelessly adapted from code by Phil Blunsom and Trevor Cohn

#include <boost/functional/hash.hpp>
#include <tr1/unordered_map>

#include "prob.h"

template <typename DishType, typename Hash = boost::hash<DishType> >
class CRP {
 public:
  CRP(double alpha) : alpha_(alpha), palpha_(alpha), total_customers_() {}
  void increment(const DishType& dish);
  void decrement(const DishType& dish);
  void erase(const DishType& dish) {
    counts_.erase(dish);
  }
  inline int count(const DishType& dish) const {
    const typename MapType::const_iterator i = counts_.find(dish);
    if (i == counts_.end()) return 0; else return i->second;
  }
  inline prob_t prob(const DishType& dish, const prob_t& p0) const {
    return (prob_t(count(dish)) + palpha_ * p0) / prob_t(total_customers_ + alpha_);
  }
 private:
  typedef std::tr1::unordered_map<DishType, int, Hash> MapType;
  MapType counts_;
  const double alpha_;
  const prob_t palpha_;
  int total_customers_;
};

template <typename Dish, typename Hash>
void CRP<Dish,Hash>::increment(const Dish& dish) {
  ++counts_[dish];
  ++total_customers_;
}

template <typename Dish, typename Hash>
void CRP<Dish,Hash>::decrement(const Dish& dish) {
  typename MapType::iterator i = counts_.find(dish);
  assert(i != counts_.end());
  if (--i->second == 0)
    counts_.erase(i);
  --total_customers_;
}

#endif