1 files changed, 48 insertions, 2 deletions

diff --git a/model/rankaccu.cpp b/model/rankaccu.cpp
index 2e77eb6..069e245 100644
--- a/model/rankaccu.cpp
+++ b/model/rankaccu.cpp
@@ -32,10 +32,49 @@ void ranksort(int l,int r,vector<int> &rank,const vector<double> &ref1,const vec
         ranksort(i,r,rank,ref1,ref2);
 }
 
+void rankmerge(int l,int r,vector<int> &C,vector<int> &rank,const vector<double> &ref1,const vector<double> &ref2)
+{
+    if (l==r)
+        return;
+    int i=l,j=((l+r)>>1)+1;
+    rankmerge(i,j-1,C,rank,ref1,ref2);
+    rankmerge(j,r,C,rank,ref1,ref2);
+    vector<int> stage_r(r-l+1),stage_c(r-l+1);
+    int cnt=0;
+    int k=0;
+    while (i<=((l+r)>>1) && j<=r)
+    {
+        if (ref1[rank[i]]>ref1[rank[j]] || (ref1[rank[i]]==ref1[rank[j]] && ref2[rank[i]]>ref2[rank[j]])) {
+            stage_c[k] = C[i];
+            stage_r[k++] = rank[i++];
+            ++cnt;
+        }
+        else{
+            stage_c[k] = C[j]+cnt;
+            stage_r[k++] = rank[j++];
+        }
+    }
+    while (i<=((l+r)>>1))
+    {
+        stage_c[k] = C[i];
+        stage_r[k++] = rank[i++];
+    }
+    while (j<=r)
+    {
+        stage_c[k] = C[j]+cnt;
+        stage_r[k++] = rank[j++];
+    }
+    for (i=l;i<=r;++i)
+    {
+        C[i]=stage_c[i-l];
+        rank[i]=stage_r[i-l];
+    }
+}
+
 int rank_accu(DataList &D,const vector<double> pred)
 {
     unsigned long n = D.getSize();
-    vector<int> orig_rank(n),pred_rank(n);
+    vector<int> orig_rank(n),pred_rank(n),C(n);
     vector<double> orig(n);
     int i,j;
     for (i=0;i<D.getSize();++i)
@@ -46,12 +85,14 @@ int rank_accu(DataList &D,const vector<double> pred)
     }
     int cnt=0;
     double accu_nDCG=0;
+    double accu_AP=0;
     i=j=0;
     while (i<D.getSize())
     {
         if ((i+1 == D.getSize())|| D.getData()[i]->qid!=D.getData()[i+1]->qid)
         {
             double Y=0,Z=0;
+            double AP=0;
             ranksort(j,i,orig_rank,orig,pred);
             ranksort(j,i,pred_rank,pred,orig);
             for (int k = j;k<=i;++k)
@@ -60,10 +101,15 @@ int rank_accu(DataList &D,const vector<double> pred)
                 Y += (pow(2,offset+orig[pred_rank[k]]) - 1)/log2(2+k-j);
             }
             accu_nDCG+=Y/Z;
+            rankmerge(j,i,C,orig_rank,pred,orig);
+            for (int k = j+1;k<=i;++k)
+                AP += ((double)C[k])/(k-j);
+            AP=AP*2/(i-j)-1;
+            accu_AP+=AP;
             j = i+1;
             ++cnt;
         }
         ++i;
     }
-    LOG(INFO)<<"Average nDGC over "<< cnt<< " queries: "<< accu_nDCG/cnt;
+    LOG(INFO)<<"over "<< cnt<< " queries. "<<"Average nDGC: "<< accu_nDCG/cnt<< " Average AP: "<<accu_AP/cnt;
 }
\ No newline at end of file