Cod sursa(job #3266777)

Utilizator lucky1992Ion Ion lucky1992 Data 10 ianuarie 2025 13:39:00
Problema Arbori indexati binar Scor 0
Compilator java Status done
Runda Arhiva educationala Marime 2.38 kb
import java.io.*;
import java.util.StringTokenizer;

public class Main {

  public static class BIT {

    private int[] tree;
    private final int maxBit;
    private int max;

    public BIT(int max) {
      this.max = max;
      tree = new int[max+1];

      int n = 1;
      while (n <= max) {
        n = n << 1;
      }
      maxBit = n;
    }

    public void update(int idx, int val) {
      while (idx <= max) {
        tree[idx] += val;
        idx += (idx & -idx);
      }
    }

    public int query(int idx) {
      int s = 0;
      while (idx > 0) {
        s += tree[idx];
        idx -= (idx & -idx);
      }
      return s;
    }

    public int findIdxForSum(int sum) {
      int bitMask = maxBit;
      int idx = 0;

      while (bitMask > 0) {
        int nextIdx = idx + bitMask;
        bitMask = bitMask >>> 1;

        if (nextIdx > max) {
          continue;
        }

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

      if (sum != 0) {
        return -1;
      } else {
        return idx;
      }
    }
  }

  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());

      BIT bit = new BIT(N);

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

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