Cod sursa(job #3266779)

Utilizator lucky1992Ion Ion lucky1992 Data 10 ianuarie 2025 14:19:05
Problema Arbori indexati binar Scor 0
Compilator java Status done
Runda Arhiva educationala Marime 2.23 kb
import java.io.*;
import java.util.StringTokenizer;

public class Main {

  public static void update(int[] bit, int i, int N, int val) {
    while (i <= N) {
      bit[i] += val;
      i += (i & -i);
    }
  }

  public static int query(int[] bit, int i) {
    int s = 0;
    while (i > 0) {
      s += bit[i];
      i -= (i & -i);
    }

    return s;
  }

  public static int search(int [] bit, int sum, int maxBit, int N) {
    int idx = 0;

    while (idx <= N && maxBit > 0) {
      int nextIdx = idx + maxBit;
      maxBit = maxBit >>> 1;

      if (nextIdx > N) {
        continue;
      }

      if (bit[nextIdx] == sum) {
        return nextIdx;
      } else if (sum > bit[nextIdx]) {
        sum -= bit[nextIdx];
        idx = nextIdx;
        if (sum == 0) {
          return idx;
        }
      }
    }

    return -1;
  }

  public static void main(String[] args) throws IOException {
    try (BufferedReader reader = new BufferedReader(new FileReader("aib.in"));
         PrintStream ps = new PrintStream(new FileOutputStream("aib.out"), true)) {

      StringTokenizer st = new StringTokenizer(reader.readLine());

      int N = Integer.parseInt(st.nextToken());
      int M = Integer.parseInt(st.nextToken());

      int[] tree = new int[N+1];

      int maxBit = 1;
      while ( (maxBit << 1) <= N) {
        maxBit = maxBit << 1;
      }

      st = new StringTokenizer(reader.readLine());
      for (int i = 1; i <= N; i++) {
        update(tree, i, N, Integer.parseInt(st.nextToken()));
      }

      for (int i = 1; i <= M; i++) {
        st = new StringTokenizer(reader.readLine());
        switch (st.nextToken()) {
          case "0":
            update(tree, Integer.parseInt(st.nextToken()), N, Integer.parseInt(st.nextToken()));
            break;
          case "1":
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            ps.println(query(tree, b) - query(tree, a-1));
            //ps.println(bit.query(b) - bit.query(a-1));
            break;
          case "2":
            int k = Integer.parseInt(st.nextToken());
            ps.println(search(tree, k, maxBit, N));
            break;
          default:
            throw new IllegalStateException();
        }
      }
    }
  }
}