fix integer overflow issue when calculating imbalance_factor (facebookresearch#4245)

Summary:

When number of clustering embeddings > int32 max, calculating imbalance factor from python side causes an function overload not found error. 
```
[0]:[rank0]:     return faiss.imbalance_factor(len(assign), k, faiss.swig_ptr(assign))
[0]:[rank0]: NotImplementedError: Wrong number or type of arguments for overloaded function 'imbalance_factor'.
[0]:[rank0]:   Possible C/C++ prototypes are:
[0]:[rank0]:     faiss::imbalance_factor(int,int,int64_t const *)
[0]:[rank0]:     faiss::imbalance_factor(int,int const *)
```

Fixing it by changing the function signature in c++ land to support int64_t.

Differential Revision: D71130612

Author: Mengdi Lin

faiss/Clustering.cpp

@@ -33,22 +33,6 @@ Clustering::Clustering(int d, int k) : d(d), k(k) {}
 Clustering::Clustering(int d, int k, const ClusteringParameters& cp)
         : ClusteringParameters(cp), d(d), k(k) {}
 
-static double imbalance_factor(int n, int k, int64_t* assign) {
-    std::vector<int> hist(k, 0);
-    for (int i = 0; i < n; i++)
-        hist[assign[i]]++;
-
-    double tot = 0, uf = 0;
-
-    for (int i = 0; i < k; i++) {
-        tot += hist[i];
-        uf += hist[i] * (double)hist[i];
-    }
-    uf = uf * k / (tot * tot);
-
-    return uf;
-}
-
 void Clustering::post_process_centroids() {
     if (spherical) {
         fvec_renorm_L2(d, k, centroids.data());

faiss/invlists/InvertedLists.cpp

@@ -181,7 +181,7 @@ size_t InvertedLists::copy_subset_to(
 }
 
 double InvertedLists::imbalance_factor() const {
-    std::vector<int> hist(nlist);
+    std::vector<int64_t> hist(nlist);
 
     for (size_t i = 0; i < nlist; i++) {
         hist[i] = list_size(i);

faiss/utils/utils.cpp

@@ -387,7 +387,7 @@ size_t ranklist_intersection_size(
     return count;
 }
 
-double imbalance_factor(int k, const int* hist) {
+double imbalance_factor(int k, const int64_t* hist) {
     double tot = 0, uf = 0;
 
     for (int i = 0; i < k; i++) {
@@ -399,9 +399,9 @@ double imbalance_factor(int k, const int* hist) {
     return uf;
 }
 
-double imbalance_factor(int n, int k, const int64_t* assign) {
-    std::vector<int> hist(k, 0);
-    for (int i = 0; i < n; i++) {
+double imbalance_factor(int64_t n, int k, const int64_t* assign) {
+    std::vector<int64_t> hist(k, 0);
+    for (int64_t i = 0; i < n; i++) {
         hist[assign[i]]++;
     }
 

faiss/utils/utils.h

@@ -92,10 +92,10 @@ size_t merge_result_table_with(
 
 /// a balanced assignment has a IF of 1, a completely unbalanced assignment has
 /// an IF = k.
-double imbalance_factor(int n, int k, const int64_t* assign);
+double imbalance_factor(int64_t n, int k, const int64_t* assign);
 
 /// same, takes a histogram as input
-double imbalance_factor(int k, const int* hist);
+double imbalance_factor(int k, const int64_t* hist);
 
 /// compute histogram on v
 int ivec_hist(size_t n, const int* v, int vmax, int* hist);