web-dev-qa-db-fra.com

Comment trouver l'index d'un élément dans un TreeSet?

J'utilise un TreeSet<Integer> et j'aimerais tout simplement trouver l'index d'un nombre dans l'ensemble. Y a-t-il une bonne façon de faire qui utilise réellement la complexité O(log(n)) des arbres binaires?

(Si non, que dois-je faire, et est-ce que quelqu'un sait pourquoi? Je suis curieux de savoir pourquoi une telle classe serait incluse dans Java sans une fonction de recherche.)

javaalgorithmdata-structurestreesetsortedset
25
jtbandes

Comme @Yrlec le fait remarquer, set.headSet(element).size retournera 0 bien qu'il n'y ait pas cet élément dans l'ensemble. Donc, nous ferions mieux de vérifier:

 return set.contains(element)? set.headSet(element).size(): -1;

Voici un cas de test pour montrer le problème:

public static void main(String args[]){
    TreeSet<Integer> set = new TreeSet<>();
    set.add(4);
    set.add(2);
    set.add(3);
    set.add(1);

    System.out.println(set.headSet(1).size());//0
    System.out.println(set.headSet(2).size());//1
    System.out.println(set.headSet(3).size());//2
    System.out.println(set.headSet(4).size());//3
    System.out.println(set.headSet(-1).size());//0!!Caution!,retusn 0 though it does not exits

}
13
Jaskey

Je me suis intéressé à TreeSet et à ses interfaces pendant un moment, et le meilleur moyen que j'ai trouvé pour obtenir l'index d'un élément est:

set.headSet(element).size()

headSet(element) renvoie la sous -TreeSet d'éléments inférieurs à son argument. La taille de cet ensemble sera donc l'index de l'élément en question. Une solution étrange en effet.

47
jtbandes

J'ai eu le même problème. J'ai donc pris le code source de Java.util.TreeMap et écrit IndexedTreeMap. Il implémente mon propre IndexedNavigableMap:

public interface IndexedNavigableMap<K, V> extends NavigableMap<K, V> {
   K exactKey(int index);
   Entry<K, V> exactEntry(int index);
   int keyIndex(K k);
}

L'implémentation est basée sur la mise à jour des poids de nœud dans l'arborescence rouge-noir lors de sa modification. Weight est le nombre de nœuds enfants situés sous un nœud donné, plus un - self. Par exemple, lorsque vous faites pivoter un arbre vers la gauche:

    private void rotateLeft(Entry<K, V> p) {
    if (p != null) {
        Entry<K, V> r = p.right;

        int delta = getWeight(r.left) - getWeight(p.right);
        p.right = r.left;
        p.updateWeight(delta);

        if (r.left != null) {
            r.left.parent = p;
        }

        r.parent = p.parent;


        if (p.parent == null) {
            root = r;
        } else if (p.parent.left == p) {
            delta = getWeight(r) - getWeight(p.parent.left);
            p.parent.left = r;
            p.parent.updateWeight(delta);
        } else {
            delta = getWeight(r) - getWeight(p.parent.right);
            p.parent.right = r;
            p.parent.updateWeight(delta);
        }

        delta = getWeight(p) - getWeight(r.left);
        r.left = p;
        r.updateWeight(delta);

        p.parent = r;
    }
  }

updateWeight met simplement à jour les poids à la racine:

   void updateWeight(int delta) {
        weight += delta;
        Entry<K, V> p = parent;
        while (p != null) {
            p.weight += delta;
            p = p.parent;
        }
    }

Et lorsque nous avons besoin de trouver l'élément par index, voici l'implémentation qui utilise des poids:

public K exactKey(int index) {
    if (index < 0 || index > size() - 1) {
        throw new ArrayIndexOutOfBoundsException();
    }
    return getExactKey(root, index);
}

private K getExactKey(Entry<K, V> e, int index) {
    if (e.left == null && index == 0) {
        return e.key;
    }
    if (e.left == null && e.right == null) {
        return e.key;
    }
    if (e.left != null && e.left.weight > index) {
        return getExactKey(e.left, index);
    }
    if (e.left != null && e.left.weight == index) {
        return e.key;
    }
    return getExactKey(e.right, index - (e.left == null ? 0 : e.left.weight) - 1);
}

Il est également très utile de trouver l’index d’une clé:

    public int keyIndex(K key) {
    if (key == null) {
        throw new NullPointerException();
    }
    Entry<K, V> e = getEntry(key);
    if (e == null) {
        throw new NullPointerException();
    }
    if (e == root) {
        return getWeight(e) - getWeight(e.right) - 1;//index to return
    }
    int index = 0;
    int cmp;
    if (e.left != null) {
        index += getWeight(e.left);
    }
    Entry<K, V> p = e.parent;
    // split comparator and comparable paths
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        while (p != null) {
            cmp = cpr.compare(key, p.key);
            if (cmp > 0) {
                index += getWeight(p.left) + 1;
            }
            p = p.parent;
        }
    } else {
        Comparable<? super K> k = (Comparable<? super K>) key;
        while (p != null) {
            if (k.compareTo(p.key) > 0) {
                index += getWeight(p.left) + 1;
            }
            p = p.parent;
        }
    }
    return index;
}

Je vais bientôt implémenter IndexedTreeSet, en attendant, vous pouvez utiliser le jeu de clés de IndexedTreeMap.

Update: IndexedTreeSet est maintenant implémenté.

Vous pouvez trouver le résultat de ce travail à l'adresse _ { http://code.google.com/p/indexed-tree-map/

12
Vitaly Sazanovich

La classe TreeSet en Java n'a pas la capacité de trouver l'index d'un nombre dans l'ensemble. Pour cela, vous devez fournir votre propre implémentation - il s'agit d'un arbre rouge-noir sous le capot, et il peut être augmenté pour prendre en charge l'opération d'index. Regardez la procédure OS-RANK dans le chapitre "Augmenter les structures de données" de "Introduction aux algorithmes", c'est précisément ce que vous demandez.

0
Óscar López