<commit_before>package de.jungblut.datastructure;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.PriorityQueue;

import com.google.common.base.Strings;
import com.google.common.collect.AbstractIterator;

import de.jungblut.distance.DistanceMeasurer;
import de.jungblut.math.DoubleVector;

/**
 * Implementation of a kd-tree.
 * 
 * @author thomas.jungblut
 * 
 */
public final class KDTree implements Iterable<DoubleVector> {

  private KDTreeNode root;

  static final class KDTreeNode {
    final int splitDimension;
    KDTreeNode left;
    KDTreeNode right;

    DoubleVector value;

    public KDTreeNode(int splitDimension, DoubleVector value) {
      this.splitDimension = splitDimension;
      this.value = value;
    }
  }

  // descending sorted by distance, so the head of the prio queue is always
  // largest
  public static final class VectorDistanceTuple implements
      Comparable<VectorDistanceTuple> {

    final DoubleVector value;
    final double dist;

    public VectorDistanceTuple(DoubleVector value, double dist) {
      this.value = value;
      this.dist = dist;
    }

    public double getDistance() {
      return dist;
    }

    public DoubleVector getVector() {
      return value;
    }

    @Override
    public int compareTo(VectorDistanceTuple o) {
      return Double.compare(o.dist, dist);
    }
  }

  /**
   * Adds the given vector to this KD tree.
   * 
   * @param vec
   */
  public void add(DoubleVector vec) {
    if (root != null) {
      KDTreeNode current = root;
      boolean right = false;
      // traverse the tree to the free spot in dimension
      while (true) {
        right = current.value.get(current.splitDimension) > vec
            .get(current.splitDimension);
        KDTreeNode next = right ? current.right : current.left;
        if (next == null) {
          break;
        } else {
          current = next;
        }
      }
      // do the real insert
      // note that current in this case is the parent
      if (right) {
        current.right = new KDTreeNode(median(vec), vec);
      } else {
        current.left = new KDTreeNode(median(vec), vec);
      }

    } else {
      root = new KDTreeNode(median(vec), vec);
    }
  }

  /**
   * @return the k nearest neighbours to the given vector.
   */
  public List<VectorDistanceTuple> getNearestNeighbours(DoubleVector v, int k,
      DistanceMeasurer measurer) {
    PriorityQueue<VectorDistanceTuple> queue = new PriorityQueue<>(k);
    searchInternal(v, root, queue, k, measurer);
    return new ArrayList<>(queue);
  }

  // TODO make more sophisticated by hyperplane intersection
  private void searchInternal(DoubleVector v, KDTreeNode current,
      PriorityQueue<VectorDistanceTuple> queue, int k, DistanceMeasurer measurer) {
    if (current != null) {
      queue.add(new VectorDistanceTuple(current.value, measurer
          .measureDistance(current.value, v)));
      if (queue.size() > k) {
        queue.remove();
      }
      boolean right = current.value.get(current.splitDimension) > v
          .get(current.splitDimension);
      KDTreeNode next = right ? current.right : current.left;
      searchInternal(v, next, queue, k, measurer);
    }
  }

  // basic in order traversal
  @Override
  public Iterator<DoubleVector> iterator() {
    return new AbstractIterator<DoubleVector>() {
      KDTreeNode current;

      @Override
      protected DoubleVector computeNext() {
        if (current == null) {
          current = root;
        } else {
          if (current.left != null) {
            current = current.left;
          }
          if (current.right != null) {
            current = current.right;
          }
          if (current.right == null && current.left == null) {
            return endOfData();
          }
        }

        return current.value;
      }

    };
  }

  @Override<commit_msg>start of the remove fnc<commit_after>package de.jungblut.datastructure;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.PriorityQueue;

import com.google.common.base.Strings;
import com.google.common.collect.AbstractIterator;

import de.jungblut.distance.DistanceMeasurer;
import de.jungblut.math.DoubleVector;

/**
 * Implementation of a kd-tree.
 * 
 * @author thomas.jungblut
 * 
 */
public final class KDTree implements Iterable<DoubleVector> {

  private KDTreeNode root;

  static final class KDTreeNode {
    final int splitDimension;
    KDTreeNode parent;
    KDTreeNode left;
    KDTreeNode right;

    DoubleVector value;

    public KDTreeNode(int splitDimension, DoubleVector value, KDTreeNode parent) {
      this.splitDimension = splitDimension;
      this.value = value;
      this.parent = parent;
    }
  }

  // descending sorted by distance, so the head of the prio queue is always
  // largest
  public static final class VectorDistanceTuple implements
      Comparable<VectorDistanceTuple> {

    final DoubleVector value;
    final double dist;

    public VectorDistanceTuple(DoubleVector value, double dist) {
      this.value = value;
      this.dist = dist;
    }

    public double getDistance() {
      return dist;
    }

    public DoubleVector getVector() {
      return value;
    }

    @Override
    public int compareTo(VectorDistanceTuple o) {
      return Double.compare(o.dist, dist);
    }
  }

  /**
   * Adds the given vector to this KD tree.
   * 
   * @param vec
   */
  public void add(DoubleVector vec) {
    if (root != null) {
      KDTreeNode current = root;
      boolean right = false;
      // traverse the tree to the free spot in dimension
      while (true) {
        right = current.value.get(current.splitDimension) > vec
            .get(current.splitDimension);
        KDTreeNode next = right ? current.right : current.left;
        if (next == null) {
          break;
        } else {
          current = next;
        }
      }
      // do the real insert
      // note that current in this case is the parent
      if (right) {
        current.right = new KDTreeNode(median(vec), vec, current);
      } else {
        current.left = new KDTreeNode(median(vec), vec, current);
      }

    } else {
      root = new KDTreeNode(median(vec), vec, null);
    }
  }

  /**
   * Removes this vector from the KD tree. It uses equals to determine if it is
   * the same vector like the passed one.
   * 
   * @return true if removed, false if not (found).
   */
  public boolean remove(DoubleVector vec) {
    KDTreeNode current = root;
    // traverse the tree to the free spot in dimension
    while (true) {

      if (vec.equals(current.value)) {
        // TODO remove the root
        // TODO the easiest way is to reinsert the nodes in the subtrees again,
        // however this is way to timeconsuming.
        if (current == root) {
          if (current.right != null) {

          } else if (current.left != null) {

          } else {
            // the root is alone, just clear up this reference
            root = null;
          }
        } else {
          // TODO removal somewhere
          if (current.left != null) {

          } else if (current.right != null) {

          } else {
            // leaf case
            if (current.parent.left == current) {
              current.parent.left = null;
            } else {
              current.parent.right = null;
            }
          }
        }
        return true;
      }
      boolean right = current.value.get(current.splitDimension) > vec
          .get(current.splitDimension);
      KDTreeNode next = right ? current.right : current.left;
      if (next == null) {
        break;
      } else {
        current = next;
      }
    }
    return false;
  }

  /**
   * @return the k nearest neighbours to the given vector.
   */
  public List<VectorDistanceTuple> getNearestNeighbours(DoubleVector vec,
      int k, DistanceMeasurer measurer) {
    PriorityQueue<VectorDistanceTuple> queue = new PriorityQueue<>(k);
    KDTreeNode current = root;

    queue.add(new VectorDistanceTuple(current.value, measurer.measureDistance(
        current.value, vec)));
    while (true) {
      if (queue.size() > k)
        queue.remove();
      boolean right = current.value.get(current.splitDimension) > vec
          .get(current.splitDimension);
      KDTreeNode next = right ? current.right : current.left;
      if (next == null) {
        break;
      } else {
        current = next;
        queue.add(new VectorDistanceTuple(current.value, measurer
            .measureDistance(current.value, vec)));
      }
    }

    return new ArrayList<>(queue);
  }

  // basic in order traversal
  @Override
  public Iterator<DoubleVector> iterator() {
    return new AbstractIterator<DoubleVector>() {
      KDTreeNode current;

      @Override
      protected DoubleVector computeNext() {
        if (current == null) {
          current = root;
        } else {
          if (current.left != null) {
            current = current.left;
          }
          if (current.right != null) {
            current = current.right;
          }
          if (current.right == null && current.left == null) {
            return endOfData();
          }
        }

        return current.value;
      }

    };
  }

  @Override
{'fraction_non_alphanumeric': 0.09757145927358693, 'fraction_numerical': 0.0, 'mean_word_length': 2.1917009602194786, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2481574', 'n_tokens_mistral': 2686, 'n_tokens_neox': 2588, 'n_words': 747}